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/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)