> JAX offers more than just a JIT compiler: JAX functions are also differentiable
if the downstream library is completely implemented in JAX (numba) ecosystem. Similar for Julia, except implementing fast code in Julia is natural, doesn't involve debugging 3 compilers (Cpython, Numba, Jax). Many python library is only differentiable because the 100x more effort were put in writing C/C++ backend, binding to python, and writing chain rules for foreign functions.
I would imagine Julia to be a good fit for this direction in the future!