expose Rng

google · Oct 28, 2023 · 7bddd53 · 7bddd53
1 parent d5d34f7
commit 7bddd53
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 0 deletions.
diff --git a/flax/experimental/nnx/__init__.py b/flax/experimental/nnx/__init__.py
@@ -92,6 +92,7 @@
 from .nnx.variables import Empty as Empty
 from .nnx.variables import Intermediate as Intermediate
 from .nnx.variables import Param as Param
+from .nnx.variables import Rng as Rng
 from .nnx.variables import Variable as Variable
 from .nnx.variables import VariableMetadata as VariableMetadata
 from .nnx.variables import with_metadata as with_metadata
diff --git a/flax/experimental/nnx/nnx/variables.py b/flax/experimental/nnx/nnx/variables.py
@@ -391,6 +391,8 @@ class Intermediate(Variable[A]):
 
 
 class Rng(Variable[jax.Array]):
+  tag: str
+
   def __init__(
     self,
     value: jax.Array,