Fix rnglib doctests for graph_updates=False default

With graph_updates=False, split_rngs and fork_rngs no longer mutate the
input node in-place. Doctests that relied on in-place mutation, restore_rngs,
or StateAxes in vmap now explicitly pass graph_updates=True.

Author: samanklesaria

flax/nnx/rnglib.py

@@ -1117,18 +1117,18 @@ def split_rngs(
     >>> from flax import nnx
     ...
     >>> rngs = nnx.Rngs(params=0, dropout=1)
-    >>> _ = nnx.split_rngs(rngs, splits=5)
+    >>> _ = nnx.split_rngs(rngs, splits=5, graph_updates=True)
     >>> rngs.params.key.shape, rngs.dropout.key.shape
     ((5,), (5,))
 
     >>> rngs = nnx.Rngs(params=0, dropout=1)
-    >>> _ = nnx.split_rngs(rngs, splits=(2, 5))
+    >>> _ = nnx.split_rngs(rngs, splits=(2, 5), graph_updates=True)
     >>> rngs.params.key.shape, rngs.dropout.key.shape
     ((2, 5), (2, 5))
 
 
     >>> rngs = nnx.Rngs(params=0, dropout=1)
-    >>> _ = nnx.split_rngs(rngs, splits=5, only='params')
+    >>> _ = nnx.split_rngs(rngs, splits=5, only='params', graph_updates=True)
     >>> rngs.params.key.shape, rngs.dropout.key.shape
     ((5,), ())
 
@@ -1140,11 +1140,11 @@ def split_rngs(
     ...     self.dropout = nnx.Dropout(0.5, rngs=rngs)
     ...
     >>> rngs = nnx.Rngs(params=0, dropout=1)
-    >>> _ = nnx.split_rngs(rngs, splits=5, only='params')
+    >>> _ = nnx.split_rngs(rngs, splits=5, only='params', graph_updates=True)
     ...
     >>> state_axes = nnx.StateAxes({(nnx.Param, 'params'): 0, ...: None})
     ...
-    >>> @nnx.vmap(in_axes=(state_axes,), out_axes=state_axes)
+    >>> @nnx.vmap(in_axes=(state_axes,), out_axes=state_axes, graph_updates=True)
     ... def create_model(rngs):
     ...   return Model(rngs)
     ...
@@ -1158,7 +1158,7 @@ def split_rngs(
 
     >>> rngs = nnx.Rngs(params=0, dropout=1)
     ...
-    >>> backups = nnx.split_rngs(rngs, splits=5, only='params')
+    >>> backups = nnx.split_rngs(rngs, splits=5, only='params', graph_updates=True)
     >>> model = create_model(rngs)
     >>> nnx.restore_rngs(backups)
     ...
@@ -1170,16 +1170,16 @@ def split_rngs(
 
     >>> rngs = nnx.Rngs(params=0, dropout=1)
     ...
-    >>> with nnx.split_rngs(rngs, splits=5, only='params'):
+    >>> with nnx.split_rngs(rngs, splits=5, only='params', graph_updates=True):
     ...   model = create_model(rngs)
     ...
     >>> model.dropout.rngs.key.shape
     ()
 
     >>> state_axes = nnx.StateAxes({(nnx.Param, 'params'): 0, ...: None})
     ...
-    >>> @nnx.split_rngs(splits=5, only='params')
-    ... @nnx.vmap(in_axes=(state_axes,), out_axes=state_axes)
+    >>> @nnx.split_rngs(splits=5, only='params', graph_updates=True)
+    ... @nnx.vmap(in_axes=(state_axes,), out_axes=state_axes, graph_updates=True)
     ... def create_model(rngs):
     ...   return Model(rngs)
     ...
@@ -1420,7 +1420,7 @@ def fork_rngs(
 
     >>> rngs = nnx.Rngs(params=0, dropout=1)
     ...
-    >>> backups = nnx.fork_rngs(rngs)
+    >>> backups = nnx.fork_rngs(rngs, graph_updates=True)
     >>> model = nnx.Linear(2, 3, rngs=rngs)
     >>> nnx.restore_rngs(backups)
     ...
@@ -1430,7 +1430,7 @@ def fork_rngs(
 
     >>> rngs = nnx.Rngs(params=0, dropout=1)
     ...
-    >>> with nnx.fork_rngs(rngs):
+    >>> with nnx.fork_rngs(rngs, graph_updates=True):
     ...   model = nnx.Linear(2, 3, rngs=rngs)
 
   """