flax

e108ab89 - Add tree-mode support to nnx.{while,fori}_loop

Add tree-mode support to nnx.{while,fori}_loop Adds `graph` parameter to `nnx.{while,fori}_loop`, enabling tree-mode operation when `graph=False`. Existing tests are parametrized to run with both graph and tree mode, plus new tree-mode-specific tests. PiperOrigin-RevId: 873126128