Add split and key to nnx.Rngs

Author: samanklesaria

flax/nnx/rnglib.py

@@ -112,15 +112,15 @@ def __init__(
 
     count = jnp.zeros(key.shape, dtype=jnp.uint32)
     self.tag = tag
-    self.key_ = RngKey(key, tag=tag)
+    self.base_key = RngKey(key, tag=tag)
     self.count = RngCount(count, tag=tag)
 
   def __call__(self) -> jax.Array:
     if not self.count.has_ref and not self.count._trace_state.is_valid():
       raise errors.TraceContextError(
         f'Cannot mutate {type(self).__name__} from a different trace level'
       )
-    key = random.fold_in(self.key_[...], self.count[...])
+    key = random.fold_in(self.base_key[...], self.count[...])
     self.count[...] += 1
     return key
 
@@ -325,7 +325,7 @@ class Rngs(Pytree):
   ``counter``. Every time a key is requested, the counter is incremented and the key is
   generated from the seed key and the counter by using ``jax.random.fold_in``.
 
-  To create an ``Rngs`` pass in an integer or ``jax.random.key_`` to the
+  To create an ``Rngs`` pass in an integer or ``jax.random.base_key`` to the
   constructor as a keyword argument with the name of the stream. The key will be used as the
   starting seed for the stream, and the counter will be initialized to zero. Then call the
   stream to get a key::
@@ -378,7 +378,7 @@ def __init__(
     Args:
       default: the starting seed for the ``default`` stream, defaults to None.
       **rngs: keyword arguments specifying the starting seed for each stream.
-        The key can be an integer or a ``jax.random.key_``.
+        The key can be an integer or a ``jax.random.base_key``.
     """
     if default is not None:
       if isinstance(default, tp.Mapping):
@@ -388,7 +388,7 @@ def __init__(
 
     for tag, key in rngs.items():
       if isinstance(key, RngStream):
-        key = key.key_[...]
+        key = key.base_key[...]
       stream = RngStream(
         key=key,
         tag=tag,
@@ -415,6 +415,9 @@ def __getattr__(self, name: str):
   def __call__(self):
     return self.default()
 
+  def key(self):
+    return self.default()
+
   def __iter__(self) -> tp.Iterator[str]:
     for name, stream in vars(self).items():
       if isinstance(stream, RngStream):
@@ -433,6 +436,9 @@ def items(self):
       if isinstance(stream, RngStream):
         yield name, stream
 
+  def split(self, splits: int):
+    return self.fork(split=splits)
+
   def fork(
     self,
     /,
@@ -457,8 +463,8 @@ def fork(
       >>> rngs = nnx.Rngs(params=1, dropout=2)
       >>> new_rngs = rngs.fork(split=5)
       ...
-      >>> assert new_rngs.params.key_.shape == (5,)
-      >>> assert new_rngs.dropout.key_.shape == (5,)
+      >>> assert new_rngs.params.base_key.shape == (5,)
+      >>> assert new_rngs.dropout.base_key.shape == (5,)
 
     ``split`` also accepts a mapping of
     `Filters <https://flax.readthedocs.io/en/latest/guides/filters_guide.html>`__  to
@@ -471,9 +477,9 @@ def fork(
       ...  ...: (2, 5),      # split anything else into 2x5 keys
       ... })
       ...
-      >>> assert new_rngs.params.key_.shape == (5,)
-      >>> assert new_rngs.dropout.key_.shape == ()
-      >>> assert new_rngs.noise.key_.shape == (2, 5)
+      >>> assert new_rngs.params.base_key.shape == (5,)
+      >>> assert new_rngs.dropout.base_key.shape == ()
+      >>> assert new_rngs.noise.base_key.shape == (2, 5)
     """
     if split is None:
       split = {}
@@ -734,18 +740,18 @@ def split_rngs(
     ...
     >>> rngs = nnx.Rngs(params=0, dropout=1)
     >>> _ = nnx.split_rngs(rngs, splits=5)
-    >>> rngs.params.key_.shape, rngs.dropout.key_.shape
+    >>> rngs.params.base_key.shape, rngs.dropout.base_key.shape
     ((5,), (5,))
 
     >>> rngs = nnx.Rngs(params=0, dropout=1)
     >>> _ = nnx.split_rngs(rngs, splits=(2, 5))
-    >>> rngs.params.key_.shape, rngs.dropout.key_.shape
+    >>> rngs.params.base_key.shape, rngs.dropout.base_key.shape
     ((2, 5), (2, 5))
 
 
     >>> rngs = nnx.Rngs(params=0, dropout=1)
     >>> _ = nnx.split_rngs(rngs, splits=5, only='params')
-    >>> rngs.params.key_.shape, rngs.dropout.key_.shape
+    >>> rngs.params.base_key.shape, rngs.dropout.base_key.shape
     ((5,), ())
 
   Once split, random state can be used with transforms like :func:`nnx.vmap`::
@@ -765,7 +771,7 @@ def split_rngs(
     ...   return Model(rngs)
     ...
     >>> model = create_model(rngs)
-    >>> model.dropout.rngs.key_.shape
+    >>> model.dropout.rngs.base_key.shape
     ()
 
   ``split_rngs`` returns a SplitBackups object that can be used to restore the
@@ -778,7 +784,7 @@ def split_rngs(
     >>> model = create_model(rngs)
     >>> nnx.restore_rngs(backups)
     ...
-    >>> model.dropout.rngs.key_.shape
+    >>> model.dropout.rngs.base_key.shape
     ()
 
   SplitBackups can also be used as a context manager to automatically restore
@@ -789,7 +795,7 @@ def split_rngs(
     >>> with nnx.split_rngs(rngs, splits=5, only='params'):
     ...   model = create_model(rngs)
     ...
-    >>> model.dropout.rngs.key_.shape
+    >>> model.dropout.rngs.base_key.shape
     ()
 
     >>> state_axes = nnx.StateAxes({(nnx.Param, 'params'): 0, ...: None})
@@ -801,7 +807,7 @@ def split_rngs(
     ...
     >>> rngs = nnx.Rngs(params=0, dropout=1)
     >>> model = create_model(rngs)
-    >>> model.dropout.rngs.key_.shape
+    >>> model.dropout.rngs.base_key.shape
     ()
 
 
@@ -828,18 +834,18 @@ def split_rngs_wrapper(*args, **kwargs):
   for path, stream in graph.iter_graph(node):
     if (
       isinstance(stream, RngStream)
-      and predicate((*path, 'key'), stream.key_)
+      and predicate((*path, 'key'), stream.base_key)
       and predicate((*path, 'count'), stream.count)
     ):
       key = stream()
-      backups.append((stream, stream.key_.raw_value, stream.count.raw_value))
+      backups.append((stream, stream.base_key.raw_value, stream.count.raw_value))
       key = random.split(key, splits)
       if squeeze:
         key = key[0]
-      if variablelib.is_array_ref(stream.key_.raw_value):
-        stream.key_.raw_value = variablelib.new_ref(key)  # type: ignore[assignment]
+      if variablelib.is_array_ref(stream.base_key.raw_value):
+        stream.base_key.raw_value = variablelib.new_ref(key)  # type: ignore[assignment]
       else:
-        stream.key_.value = key
+        stream.base_key.value = key
       if squeeze:
         counts_shape = stream.count.shape
       elif isinstance(splits, int):
@@ -898,18 +904,18 @@ def fork_rngs(
     ...
     >>> rngs = nnx.Rngs(params=0, dropout=1)
     >>> _ = nnx.fork_rngs(rngs, split=5)
-    >>> rngs.params.key_.shape, rngs.dropout.key_.shape
+    >>> rngs.params.base_key.shape, rngs.dropout.base_key.shape
     ((5,), (5,))
 
     >>> rngs = nnx.Rngs(params=0, dropout=1)
     >>> _ = nnx.fork_rngs(rngs, split=(2, 5))
-    >>> rngs.params.key_.shape, rngs.dropout.key_.shape
+    >>> rngs.params.base_key.shape, rngs.dropout.base_key.shape
     ((2, 5), (2, 5))
 
 
     >>> rngs = nnx.Rngs(params=0, dropout=1)
     >>> _ = nnx.fork_rngs(rngs, split={'params': 5})
-    >>> rngs.params.key_.shape, rngs.dropout.key_.shape
+    >>> rngs.params.base_key.shape, rngs.dropout.base_key.shape
     ((5,), ())
 
   Once forked, random state can be used with transforms like :func:`nnx.vmap`::
@@ -929,7 +935,7 @@ def fork_rngs(
     ...   return Model(rngs)
     ...
     >>> model = create_model(rngs)
-    >>> model.dropout.rngs.key_.shape
+    >>> model.dropout.rngs.base_key.shape
     ()
 
   ``fork_rngs`` returns a SplitBackups object that can be used to restore the
@@ -942,7 +948,7 @@ def fork_rngs(
     >>> model = create_model(rngs)
     >>> nnx.restore_rngs(backups)
     ...
-    >>> model.dropout.rngs.key_.shape
+    >>> model.dropout.rngs.base_key.shape
     ()
 
   SplitBackups can also be used as a context manager to automatically restore
@@ -953,7 +959,7 @@ def fork_rngs(
     >>> with nnx.fork_rngs(rngs, split={'params': 5}):
     ...   model = create_model(rngs)
     ...
-    >>> model.dropout.rngs.key_.shape
+    >>> model.dropout.rngs.base_key.shape
     ()
 
     >>> state_axes = nnx.StateAxes({(nnx.Param, 'params'): 0, ...: None})
@@ -965,7 +971,7 @@ def fork_rngs(
     ...
     >>> rngs = nnx.Rngs(params=0, dropout=1)
     >>> model = create_model(rngs)
-    >>> model.dropout.rngs.key_.shape
+    >>> model.dropout.rngs.base_key.shape
     ()
   """
   if isinstance(node, Missing):
@@ -993,14 +999,14 @@ def fork_rngs_wrapper(*args, **kwargs):
     for predicate, splits in predicate_splits.items():
       if (
         isinstance(stream, RngStream)
-        and predicate((*path, 'key'), stream.key_)
+        and predicate((*path, 'key'), stream.base_key)
         and predicate((*path, 'count'), stream.count)
       ):
         forked_stream = stream.fork(split=splits)
         # backup the original stream state
-        backups.append((stream, stream.key_.raw_value, stream.count.raw_value))
+        backups.append((stream, stream.base_key.raw_value, stream.count.raw_value))
         # apply the forked key and count to the original stream
-        stream.key_.raw_value = forked_stream.key_.raw_value
+        stream.base_key.raw_value = forked_stream.base_key.raw_value
         stream.count.raw_value = forked_stream.count.raw_value
 
   return SplitBackups(backups)
@@ -1010,7 +1016,7 @@ def backup_keys(node: tp.Any, /):
   backups: list[StreamBackup] = []
   for _, stream in graph.iter_graph(node):
     if isinstance(stream, RngStream):
-      backups.append((stream, stream.key_.raw_value))
+      backups.append((stream, stream.base_key.raw_value))
   return backups
 
 def _scalars_only(
@@ -1055,7 +1061,7 @@ def reseed(
       of the form ``(path, scalar_key, target_shape) -> new_key`` can be passed to
       define a custom reseeding policy.
     **stream_keys: a mapping of stream names to new keys. The keys can be
-      either integers or ``jax.random.key_``.
+      either integers or ``jax.random.base_key``.
 
   Example::
 
@@ -1093,16 +1099,16 @@ def reseed(
   rngs = Rngs(**stream_keys)
   for path, stream in graph.iter_graph(node):
     if isinstance(stream, RngStream):
-      if stream.key_.tag in stream_keys:
-        key = rngs[stream.key_.tag]()
-        key = policy(path, key, stream.key_.shape)
-        stream.key_.value = key
+      if stream.base_key.tag in stream_keys:
+        key = rngs[stream.base_key.tag]()
+        key = policy(path, key, stream.base_key.shape)
+        stream.base_key.value = key
         stream.count.value = jnp.zeros(key.shape, dtype=jnp.uint32)
 
 
 def restore_rngs(backups: tp.Iterable[StreamBackup], /):
   for backup in backups:
     stream = backup[0]
-    stream.key_.raw_value = backup[1]
+    stream.base_key.raw_value = backup[1]
     if len(backup) == 3:
       stream.count.raw_value = backup[2]  # count

tests/nnx/bridge/module_test.py

@@ -149,7 +149,7 @@ def __call__(self):
     scope = bar.apply({}, rngs=1)
     self.assertIsNone(bar.scope)
 
-    self.assertEqual(scope.rngs.default.key_[...], jax.random.key(1))
+    self.assertEqual(scope.rngs.default.base_key[...], jax.random.key(1))
     self.assertEqual(scope.rngs.default.count[...], 0)
 
     class Baz(bridge.Module):

tests/nnx/module_test.py

@@ -556,7 +556,7 @@ def test_create_abstract(self):
   def test_create_abstract_stateful(self):
     linear = nnx.eval_shape(lambda: nnx.Dropout(0.5, rngs=nnx.Rngs(0)))
 
-    assert linear.rngs.key_.value == jax.ShapeDtypeStruct(
+    assert linear.rngs.base_key.value == jax.ShapeDtypeStruct(
       (), jax.random.key(0).dtype
     )
 

tests/nnx/mutable_array_test.py

@@ -662,7 +662,7 @@ def test_rngs_create(self):
       paths[1],
       (
         jax.tree_util.GetAttrKey('default'),
-        jax.tree_util.GetAttrKey('key_'),
+        jax.tree_util.GetAttrKey('base_key'),
         jax.tree_util.GetAttrKey('value'),
       ),
     )

tests/nnx/nn/attention_test.py

@@ -128,7 +128,7 @@ def test_keep_rngs(self, keep_rngs):
     if keep_rngs:
       _, _, nondiff = nnx.split(module, nnx.Param, ...)
       assert isinstance(nondiff['rngs']['count'], nnx.RngCount)
-      assert isinstance(nondiff['rngs']['key_'], nnx.RngKey)
+      assert isinstance(nondiff['rngs']['base_key'], nnx.RngKey)
     else:
       nnx.split(module, nnx.Param)
 

tests/nnx/rngs_test.py

@@ -45,12 +45,12 @@ def test_rng_stream(self):
 
     key1 = rngs.params()
     self.assertEqual(rngs.params.count[...], 1)
-    self.assertIs(rngs.params.key_[...], key0)
+    self.assertIs(rngs.params.base_key[...], key0)
     self.assertFalse(jnp.allclose(key0, key1))
 
     key2 = rngs.params()
     self.assertEqual(rngs.params.count[...], 2)
-    self.assertIs(rngs.params.key_[...], key0)
+    self.assertIs(rngs.params.base_key[...], key0)
     self.assertFalse(jnp.allclose(key1, key2))
 
   def test_rng_trace_level_constraints(self):