jax does have a just in time compiler that lets you compile your python functions to XLA-optimized kernels (through llvmlite under the hood?).

The fact that you can jit compile and gain the benefits of "doing this within the compiler" is one of its main selling points.