flax

4f9f0acb - Fix update doctest for graph_updates=False default

Fix update doctest for graph_updates=False default With graph_updates=False, nnx.grad returns a graph node rather than a State, so nnx.state() must be called on grads before passing to jax.tree.map.