I’ve been disappointed with Jax which I was trying to use for backward auto differentiation. The issue is that XLA JIT compilation is very slow and easily adds

I’ve been disappointed with Jax which I was trying to use for backward auto differentiation. The issue is that XLA JIT compilation is very slow and easily adds half a minute of overhead to the first call of the base function just by using jax.numpy instead of numpy, which made it a non starter for my use case. It’s definitely optimised for large flow computations where the JIT overhead is dwarfed by the rest. In the end I reverted to using autograd which did the job fine. I had never heard of tai chi until now, I’m curious how it compares.