flax

49213e07 - Add graph=False support for nnx.grad and nnx.value_and_grad

Add graph=False support for nnx.grad and nnx.value_and_grad Add tree-mode support to nnx.grad/nnx.value_and_grad, where Variables are treated as pytree leaves and argnums are passed directly to JAX. PiperOrigin-RevId: 869925351