define separated jit and normal kernels in order to reduce compilation time
hghcomphys opened this issue · comments
As it's recommended, by the JAX's documentation, to call jax.jit on the big function.
This reduces compilation time and allowing better optimization done by the compiler.