Merge branch 'remove-variable-ref' into hijax-variable

Author: cgarciae

docs_nnx/guides/array_ref.ipynb

@@ -123,8 +123,8 @@
     }
    ],
    "source": [
-    "variable = nnx.Variable(jnp.array([1, 2, 3]), use_ref=True)\n",
-    "print(f\"{variable.has_ref = }\\n\")\n",
+    "variable = nnx.Variable(jnp.array([1, 2, 3]), is_hijax=True)\n",
+    "print(f\"{variable.is_hijax = }\\n\")\n",
     "\n",
     "print(\"[1] =\", variable); increment(variable); print(\"[2] =\", variable)"
    ]
@@ -147,7 +147,7 @@
     "with nnx.use_refs(True):\n",
     "  variable = nnx.Variable(jnp.array([1, 2, 3]))\n",
     "\n",
-    "print(f\"{variable.has_ref = }\")"
+    "print(f\"{variable.is_hijax = }\")"
    ]
   },
   {

docs_nnx/guides/array_ref.md

@@ -45,8 +45,8 @@ print(increment.lower(a_ref).as_text())
 ### Variables Refs
 
 ```{code-cell} ipython3
-variable = nnx.Variable(jnp.array([1, 2, 3]), use_ref=True)
-print(f"{variable.has_ref = }\n")
+variable = nnx.Variable(jnp.array([1, 2, 3]), is_hijax=True)
+print(f"{variable.is_hijax = }\n")
 
 print("[1] =", variable); increment(variable); print("[2] =", variable)
 ```
@@ -55,7 +55,7 @@ print("[1] =", variable); increment(variable); print("[2] =", variable)
 with nnx.use_refs(True):
   variable = nnx.Variable(jnp.array([1, 2, 3]))
 
-print(f"{variable.has_ref = }")
+print(f"{variable.is_hijax = }")
 ```
 
 Mention `nnx.use_refs` can be used as global flag

flax/nnx/__init__.py

@@ -193,8 +193,6 @@
 from .variablelib import variable_type_from_name as variable_type_from_name
 from .variablelib import variable_name_from_type as variable_name_from_type
 from .variablelib import register_variable_name as register_variable_name
-from .variablelib import use_refs as use_refs
-from .variablelib import using_refs as using_refs
 from .variablelib import use_hijax as use_hijax
 from .variablelib import using_hijax as using_hijax
 from .visualization import display as display

flax/nnx/bridge/module.py

@@ -388,10 +388,9 @@ def _get_variables(self) -> tp.Mapping:
       if collection not in _variables:
         _variables[collection] = {}
 
-      if (
-        isinstance(variable, variablelib.Variable)
-        and not variable.get_metadata()
-      ):
+      if isinstance(
+        variable, variablelib.Variable
+      ) and bridge_variables.is_vanilla_variable(variable):
         leaf = variable.value
       else:
         leaf = bridge_variables.to_linen_var(variable)

flax/nnx/bridge/variables.py

@@ -78,11 +78,11 @@ def is_vanilla_variable(vs: variablelib.Variable) -> bool:
   Returns False only if it has non-empty hooks or any non-built-in attribute.
   """
   for key, value in vs.get_metadata().items():
-    if key.endswith('_hooks'):
-      if value != ():
-        return False
-    else:
-      return False
+    if key in ('is_hijax', 'eager_sharding'):
+      continue
+    if key.endswith('_hooks') and value == ():
+      continue
+    return False
   return True
 
 

flax/nnx/graph.py

@@ -20,7 +20,7 @@
 import threading
 import typing as tp
 
-import jax.experimental
+import jax.core
 
 from flax import config
 from flax.nnx import filterlib, reprlib, traversals, variablelib
@@ -373,17 +373,13 @@ class VariableDef(reprlib.Representable, tp.Generic[Node]):
   index: int
   outer_index: int | None
   metadata: HashableMapping[str, tp.Any]
-  array_refdef: ArrayRefDef | NodeRef | None
 
   def with_no_outer_index(self) -> VariableDef:
     return VariableDef(
       type=self.type,
       index=self.index,
       outer_index=None,
       metadata=self.metadata,
-      array_refdef=self.array_refdef.with_no_outer_index()
-      if isinstance(self.array_refdef, ArrayRefDef)
-      else self.array_refdef,
     )
 
   def with_same_outer_index(self) -> VariableDef:
@@ -392,9 +388,6 @@ def with_same_outer_index(self) -> VariableDef:
       index=self.index,
       outer_index=self.index,
       metadata=self.metadata,
-      array_refdef=self.array_refdef.with_same_outer_index()
-      if isinstance(self.array_refdef, ArrayRefDef)
-      else self.array_refdef,
     )
 
   def __nnx_repr__(self):
@@ -761,32 +754,23 @@ def make_mutable_arraydef(value: variablelib.Ref):
   if is_variable:
     assert isinstance(node, Variable)
     assert index is not None
-    prev_inner_value = node.raw_value
-    if variablelib.is_array_ref(prev_inner_value):
-      array_refdef, inner_value = make_mutable_arraydef(prev_inner_value)
-    else:
-      array_refdef = None
-      inner_value = prev_inner_value
     if path is None:
-      leaf = inner_value
+      leaf = node.raw_value
     else:
       leaf = node  # type: ignore[assignment]
-      if inner_value is not prev_inner_value:
-        leaf.raw_value = inner_value
 
     variabledef = VariableDef(
-      type=type(node),
+      type=jax.typeof(node)._var_type  # type: ignore
+      if isinstance(node, jax.core.Tracer)
+      else type(node),
       index=index,
       outer_index=ref_outer_index.get(node, None) if ref_outer_index else None,
       metadata=HashableMapping(node.get_metadata()),
-      array_refdef=array_refdef,
     )
-    if type(inner_value) is not Repeated:
-      assert not isinstance(leaf, Repeated)
-      leaves.append(leaf)
-      if path is not None:
-        assert paths is not None
-        paths.append(tuple(path))
+    leaves.append(leaf)
+    if path is not None:
+      assert paths is not None
+      paths.append(tuple(path))
     nodes.append(variabledef)
     return
   elif is_array_ref:
@@ -1200,7 +1184,7 @@ def get_mutable_array(array_refdef: ArrayRefDef, leaf):
         f"Expected a ArrayRefOutput type but got '{leaf.value}.'"
       )
     elif type(leaf) is ArrayRefOutput:
-      array_ref = variablelib.new_ref(leaf.value)
+      array_ref = jax.new_ref(leaf.value)
     elif variablelib.is_array_ref(leaf):
       array_ref = leaf
     else:
@@ -1217,26 +1201,9 @@ def get_mutable_array(array_refdef: ArrayRefDef, leaf):
     variabledef = tp.cast(VariableDef[Variable], nodedef)
     # its a unseen variable, create a new one
 
-    if variabledef.array_refdef is not None:
-      if type(variabledef.array_refdef) is NodeRef:
-        value = index_ref[variabledef.array_refdef.index]
-      else:
-        value = next(leaves_iter)
-        assert type(variabledef.array_refdef) is ArrayRefDef
-        if isinstance(value, Variable):
-          value = value.copy() if copy_variables else value
-          inner_value = value.raw_value
-          array_ref = get_mutable_array(variabledef.array_refdef, inner_value)
-          if array_ref is not inner_value:
-            value.raw_value = array_ref
-        else:
-          # if value is an array or array ref, we need call get_mutable_array
-          # to register it in the index_ref
-          value = get_mutable_array(variabledef.array_refdef, value)
-    else:
-      value = next(leaves_iter)
-      if isinstance(value, Variable) and copy_variables:
-        value = value.copy()
+    value = next(leaves_iter)
+    if isinstance(value, Variable) and copy_variables:
+      value = value.copy()
 
     # when idxmap is present, check if the Varable exists there
     # and update existing variables if it does
@@ -1442,7 +1409,7 @@ def _update_variable(node: Variable, value):
         ):
           node[...] = value[...]
         else:
-          node.raw_value = value
+          node.set_raw_value(value)
 
   if isinstance(node, Variable):
     _update_variable(node, state)
@@ -2616,7 +2583,7 @@ def clone(node: Node, variables: bool = True) -> Node:
 
 
 def _mutable_like(path, x):
-  return (isinstance(x, Variable) and x.has_ref) or variablelib.is_array_ref(x)
+  return variablelib.is_array_ref(x)
 
 
 def to_arrays(
@@ -2669,7 +2636,9 @@ def to_arrays(
   Returns:
     A structure with the frozen arrays.
   """
-  if not allow_duplicates and (all_duplicates := find_duplicates(node, only=only)):
+  if not allow_duplicates and (
+    all_duplicates := find_duplicates(node, only=only)
+  ):
     duplicates_strs = '\n  ---'
     for node_duplicates in all_duplicates:
       for path in node_duplicates:
@@ -2685,7 +2654,7 @@ def to_arrays(
 
 
 def _array_like(path, x):
-  return (isinstance(x, Variable) and not x.has_ref) or isinstance(x, jax.Array)
+  return isinstance(x, jax.Array)
 
 
 def to_refs(node: A, /, only: filterlib.Filter = _array_like) -> A:
@@ -2741,13 +2710,13 @@ def to_refs(node: A, /, only: filterlib.Filter = _array_like) -> A:
     raise ValueError(f'Found duplicate at paths:{duplicates_strs}')
 
   graphdef, frozen_state, rest = split(node, only, ...)  # type: ignore[misc]
-  mutable_state = jax.tree.map(variablelib.new_ref, frozen_state)
+  mutable_state = jax.tree.map(jax.new_ref, frozen_state)
   node = merge(graphdef, mutable_state, rest)
   return node
 
 def _is_lojax_variable(path, x):
   return isinstance(x, variablelib.Variable) and not isinstance(
-    x, variablelib.MutableHijaxVariable
+    x, variablelib.HijaxVariable
   )
 
 
@@ -2789,7 +2758,7 @@ def _to_stateful(x):
 
 
 def _is_hijax_variable(path, x):
-  return isinstance(x, variablelib.MutableHijaxVariable)
+  return isinstance(x, variablelib.HijaxVariable)
 
 
 def to_lojax(node: A, /, only: filterlib.Filter = ...) -> A:
@@ -2808,7 +2777,7 @@ def to_lojax(node: A, /, only: filterlib.Filter = ...) -> A:
   def _to_stateless(x):
     if variablelib.is_array_ref(x):
       return x[...]
-    elif isinstance(x, variablelib.MutableHijaxVariable):
+    elif isinstance(x, variablelib.HijaxVariable):
       return variablelib._get_mutable_hijax_state(x)
     return x
 
@@ -2820,6 +2789,75 @@ def _to_stateless(x):
   return node
 
 
+def _is_lojax_variable(path, x):
+  return isinstance(x, variablelib.Variable) and not isinstance(
+    x, variablelib.HijaxVariable
+  )
+
+
+def to_hijax(
+  node: A, /, *, only: filterlib.Filter = ..., mutable: bool = True
+) -> A:
+  """ """
+  if not mutable:
+    raise ValueError('to_hijax only supports mutable=True at the moment.')
+
+  only = filterlib.All(_is_lojax_variable, only)
+  predicate = filterlib.to_predicate(only)
+
+  if all_duplicates := find_duplicates(node, only=only):
+    duplicates_strs = '\n  ---'
+    for node_duplicates in all_duplicates:
+      for path in node_duplicates:
+        path_str = '/'.join(map(str, path))
+        duplicates_strs += f'\n  {path_str}'
+      duplicates_strs += '\n  ---'
+    raise ValueError(f'Found duplicate at paths:{duplicates_strs}')
+
+  def _to_hijax(jax_path, x):
+    if predicate(to_nnx_path(jax_path), x):
+      assert isinstance(x, variablelib.Variable)
+      x = x.copy()
+      x._var_metadata['is_hijax'] = True
+      return variablelib._new_mutable_hijax_from_variable(x)
+    return x
+
+  node = jax.tree.map_with_path(
+    _to_hijax, node, is_leaf=lambda x: isinstance(x, variablelib.Variable)
+  )
+  return node
+
+
+def _is_hijax_variable(path, x):
+  return isinstance(x, variablelib.HijaxVariable)
+
+def to_lojax(node: A, /, only: filterlib.Filter = ...) -> A:
+  """ """
+  only = filterlib.All(_is_hijax_variable, only)
+  predicate = filterlib.to_predicate(only)
+
+  if all_duplicates := find_duplicates(node, only=only):
+    duplicates_strs = '\n  ---'
+    for node_duplicates in all_duplicates:
+      for path in node_duplicates:
+        path_str = '/'.join(map(str, path))
+        duplicates_strs += f'\n  {path_str}'
+      duplicates_strs += '\n  ---'
+    raise ValueError(f'Found duplicate at paths:{duplicates_strs}')
+
+  def _to_lojax(jax_path, x):
+    if predicate(to_nnx_path(jax_path), x):
+      variable = variablelib._get_mutable_hijax_state(x)
+      variable._var_metadata['is_hijax'] = False
+      return variable
+    return x
+
+  node = jax.tree.map_with_path(
+    _to_lojax, node, is_leaf=lambda x: isinstance(x, variablelib.Variable)
+  )
+  return node
+
+
 def pure(tree: A) -> A:
   """Returns a new tree with all ``Variable`` objects replaced with inner values.
 
@@ -3144,7 +3182,7 @@ def _key_path_to_key(key: tp.Any) -> Key:
     return str(key)
 
 
-def jax_to_nnx_path(jax_path: tuple, /):
+def to_nnx_path(jax_path: tuple, /):
   return tuple(_key_path_to_key(part) for part in jax_path)
 
 

flax/nnx/nn/normalization.py

@@ -54,7 +54,7 @@ def _compute_stats(
   use_mean: bool = True,
   use_fast_variance: bool = True,
   mask: tp.Optional[Array] = None,
-):
+) -> tuple[Array, Array]:
   """Computes mean and variance statistics.
 
   This implementation takes care of a few important details:
@@ -357,6 +357,8 @@ def __call__(
     feature_axes = _canonicalize_axes(x.ndim, self.axis)
     reduction_axes = tuple(i for i in range(x.ndim) if i not in feature_axes)
 
+    mean: jax.Array
+    var: jax.Array
     if use_running_average:
       mean, var = self.mean.value, self.var.value
     else:
@@ -370,7 +372,7 @@ def __call__(
         mask=mask,
       )
       # stop_gradient only for flax_array_ref
-      if self.mean.has_ref or self.var.has_ref:
+      if self.mean.is_hijax or self.var.is_hijax:
         stop_gradient = jax.lax.stop_gradient
       else:
         stop_gradient = lambda x: x