Ensure compatibility NNX 0.11.2

[juanitorduz]

Ensures compatibility with https://github.com/google/flax/releases/tag/v0.11.2

The tests were failing in https://github.com/pyro-ppl/numpyro/actions/runs/17463007632/job/49592472608

FAILED test/contrib/test_module.py::test_nnx_state_dropout_smoke[batchnorm-dropout] - flax.errors.TraceContextError: Cannot mutate BatchStat from a different trace level (https://flax.readthedocs.io/en/latest/api_reference/flax.errors.html#flax.errors.TraceContextError)
--------------------
For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.
FAILED test/contrib/test_module.py::test_nnx_state_dropout_smoke[batchnorm-no_dropout] - flax.errors.TraceContextError: Cannot mutate BatchStat from a different trace level (https://flax.readthedocs.io/en/latest/api_reference/flax.errors.html#flax.errors.TraceContextError)
--------------------
For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.
= 2 failed, 506 passed, 66 skipped, 4 xfailed, 36 warnings in 654.96s (0:10:54)

[fehiepsi]

Thanks for the fix @juanitorduz!

[juanitorduz]

I will rebase and remove the xfail once #2069 is merged

[fehiepsi]

Hi @juanitorduz, I just merged the other PR.

[juanitorduz]

Hi @juanitorduz, I just merged the other PR.

I just rebased and removed xfail