hijax Variable

Author: cgarciae

flax/nnx/extract.py

@@ -18,6 +18,7 @@
 import jax
 
 from flax import struct
+from flax import typing
 from flax.nnx.pytreelib import Pytree
 from flax.typing import Missing, PathParts
 from flax.nnx import graph, variablelib
@@ -35,7 +36,7 @@ class PrefixMapping(abc.ABC):
   @abc.abstractmethod
   def map_prefix(
     self,
-    path: variablelib.PathParts,
+    path: typing.PathParts,
     variable: variablelib.Variable,
     /,
   ) -> tp.Any: ...

flax/nnx/rnglib.py

@@ -70,7 +70,11 @@ def __init__(
     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():
+    if (
+      not self.count.has_ref
+      and self.count.trace_state
+      and not self.count.trace_state.is_valid()
+    ):
       raise errors.TraceContextError(
         f'Cannot mutate {type(self).__name__} from a different trace level'
       )

flax/nnx/transforms/iteration.py

@@ -19,6 +19,7 @@
 import typing as tp
 
 from flax import struct
+from flax import typing
 from flax.core.frozen_dict import FrozenDict
 from flax.nnx import extract, filterlib, graph, spmd, variablelib
 from flax.nnx import statelib
@@ -89,7 +90,7 @@ def axes(self) -> tuple[Index | type[Carry] | None, ...]:
     return self._axes
 
   def map_prefix(
-    self, path: variablelib.PathParts, variable: variablelib.Variable
+    self, path: typing.PathParts, variable: variablelib.Variable
   ) -> tp.Any:
     for filter, axis in zip(self.filters, self.axes):
       predicate = filterlib.to_predicate(filter)