Improve chex.fake_jit by also disabling internal jitting in functions such as jax.lax.scan.
Improve chex.fake_jit by also disabling internal jitting in functions such as jax.lax.scan.

PiperOrigin-RevId: 393347523

deepmind/chex,deepmind/chex