# Global flag
jax.config.update("jax_disable_jit", True)

# Or via environment variable, before starting Python
# JAX_DISABLE_JIT=1 python train.py

        __
        __