remove jax.jit decorator from small functions/methods
hghcomphys opened this issue · comments
It's recommended by JAX's documentaiton that JIT should be used for bigger functions.
This will cause unnecessary compilation each time in the loop, even if JAX uses cached functions but still relies on hashing.