What are current limitations for using jax.jit for training?

[wittenator]

For example the guide https://flax.readthedocs.io/en/latest/guides/flax_gspmd.html is completely written with jax.jit, but using it for my own more complex models seems to not update the weights during training. Replacing nnx.value_and_grad with jax.value_and_grad also throws errors due to int values. Are there known cases where this cannot be used such with randomness etc.?