Add split and fork methods to RngStream

Author: samanklesaria

flax/nnx/rnglib.py

@@ -124,6 +124,12 @@ def __call__(self) -> jax.Array:
     self.count[...] += 1
     return key
 
+  def key(self) -> jax.Array:
+    return self()
+
+  def split(self, k: int):
+      return self.fork(split=k)
+
   def fork(self, *, split: int | tuple[int, ...] | None = None):
     key = self()
     if split is not None: