Detect mixing of linen and nnx api

[Lime-Cakes]

I'm recently upgrading old nn modules to nnx and missed one. Training code reported TraceContextError.
flax.errors.TraceContextError: Cannot mutate RngCount from a different trace level

It's a fairly confusing error. Would be nice if some assertion's added to check if nn and nnx modules get mixed.

Edit: My problem's caused by mixing nn and nnx module together, causing the mutation inside jitted function. Leaving it here in case someone else ran into same problem. Don't waste hours debugging a typo.

[vfdev-5]

The error is about a mutation of rng inside jitted function. Usually, cloning the model or rng can help with it:

• Dropout can yield TraceContextError is nnx.merge(..., copy=False) #5050 (comment)
• Train step with nnx.jit/jax.jit+nnx.grad consumes twice memory than jax.jit+jax.grad equivalent #5116 (comment)

[Lime-Cakes]

Thanks! Kinda unrelated question, is it true that flax aims to migrate to jax.jit/jax.grad instead of nnx.jit/nnx.grad?

[vfdev-5]

Yes, it is preferred to use Jax transforms instead of nnx ones