allow split tuples in Rngs.fork

Author: cgarciae

flax/nnx/rnglib.py

@@ -623,6 +623,7 @@ def fork(
     *,
     split: tp.Mapping[filterlib.Filter, int | tuple[int, ...]]
     | int
+    | tuple[int, ...]
     | None = None,
   ):
     """Returns a new Rngs object with new unique RNG keys.
@@ -662,6 +663,8 @@ def fork(
       split = {}
     elif isinstance(split, int):
       split = {...: split}
+    elif isinstance(split, tuple):
+      split = {...: split}
 
     split_predicates = {filterlib.to_predicate(k): v for k, v in split.items()}
     keys: dict[str, RngStream] = {}