Add split and key to nnx.Rngs

Author: samanklesaria

flax/nnx/rnglib.py

@@ -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,
     /,