fix(nnx): Add backward-compatibility safe accessors for Flax NNX containers

Author: hsuan-lun-chiang

src/maxtext/common/train_state_nnx.py

@@ -220,8 +220,16 @@ def split_for_checkpoint(state: nnx.State):
   params, rng_state, batch_stats, caches, intermediates, rest = nnx.split_state(
       state, nnx.Param, nnx.RngState, nnx.BatchStat, nnx.Cache, nnx.Intermediate, ...
   )
-  optimizer = nnx.State({"optimizer": rest["optimizer"]}) if "optimizer" in rest else nnx.State({})
-  custom = nnx.State({"model": rest["model"]}) if "model" in rest else nnx.State({})
+  optimizer = (
+      nnx.State({"optimizer": rest["optimizer"].to_pure_dict() if hasattr(rest["optimizer"], "to_pure_dict") else rest["optimizer"]})
+      if "optimizer" in rest
+      else nnx.State({})
+  )
+  custom = (
+      nnx.State({"model": rest["model"].to_pure_dict() if hasattr(rest["model"], "to_pure_dict") else rest["model"]})
+      if "model" in rest
+      else nnx.State({})
+  )
   linen_state = nnx.merge_state(params, optimizer)
   aux = nnx.merge_state(rng_state, batch_stats, custom)
   ephemeral = nnx.merge_state(caches, intermediates)

src/maxtext/layers/embeddings.py

@@ -153,7 +153,10 @@ def __call__(self, inputs: Array, model_mode: str = MODEL_MODE_TRAIN) -> Array:
       raise ValueError("Input type must be an integer or unsigned integer.")
 
     embedding = jnp.asarray(
-        _maybe_move_embedding_to_device(self.embedding.get_value(), self.config),
+        _maybe_move_embedding_to_device(
+            self.embedding.raw_value if hasattr(self.embedding, "raw_value") else getattr(self.embedding, "value", self.embedding),
+            self.config,
+        ),
         self.dtype,
     )
 
@@ -197,7 +200,7 @@ def attend(self, query: Array, out_sharding: NamedSharding | None = None) -> Arr
       Commonly used for weight-sharing between embeddings and logit transform
       in NLP models.
     """
-    embedding = self.embedding.get_value()
+    embedding = self.embedding.raw_value if hasattr(self.embedding, "raw_value") else getattr(self.embedding, "value", self.embedding)
     attend_dtype = self.attend_dtype if self.attend_dtype is not None else self.dtype
     return attend_on_embedding(query, embedding, attend_dtype, self.config, out_sharding)
 

src/maxtext/trainers/pre_train/train.py

@@ -78,9 +78,8 @@
 def get_first_step(model, state):
   if isinstance(model, nn.Module):
     return int(state.step)
-  if hasattr(state, "inner_state"):  # DiLoCoTrainState (NNX DiLoCo): step is the optimizer step var
-    return int(state.step.get_value())
-  return int(state.optimizer.step.get_value())
+  step_var = state.step if hasattr(state, "inner_state") else state.optimizer.step
+  return int(step_var.value if hasattr(step_var, "value") else step_var.raw_value if hasattr(step_var, "raw_value") else step_var)
 
 
 # -----------------------------------------------------------------------------
@@ -420,7 +419,10 @@ def train_step(model, config, state_mesh_shardings, params_shardings, state, dat
         nnx.update(state.model, curr_params)
 
       def diff_wrapper(curr_params, custom_params, rest, config, data):
-        local_model = nnx.merge(model_graphdef, curr_params, custom_params, rest, copy=True)
+        try:
+          local_model = nnx.merge(model_graphdef, curr_params, custom_params, rest, copy=True)
+        except TypeError:
+          local_model = nnx.merge(model_graphdef, curr_params, custom_params, rest)
         loss, aux = loss_fn(local_model, config, data, None, None, is_train=True)
         _, _, _, new_rest = nnx.split(local_model, nnx.Param, custom_param_filter, ...)
         return loss, (aux, new_rest)

src/maxtext/utils/maxtext_utils_nnx.py

@@ -180,12 +180,22 @@ def create_sharded_state():
   return nnx.merge(graphdef, sharded_state)
 
 
+def _get_val(x):
+  if hasattr(x, "raw_value"):
+    return x.raw_value
+  if hasattr(x, "value"):
+    return x.value
+  if hasattr(x, "get_value"):
+    return x.get_value()
+  return x
+
+
 def nnx_ensure_scan_leading_axis(tree, length):
   """Broadcasts scalar-like variables to have a leading scan axis."""
 
   def _op(x):
     is_var = isinstance(x, nnx.Variable)
-    val = x.get_value() if is_var else x
+    val = _get_val(x)
     if hasattr(val, "shape") and len(val.shape) == 0:
       new_val = jax.numpy.broadcast_to(val, (length,))
       return x.replace(value=new_val) if is_var else new_val
@@ -214,6 +224,8 @@ def nnx_update_sharding_meta(variable, transform_fn):
       updates[key] = P(*transformed) if isinstance(val, P) else tuple(transformed)
 
   if updates:
+    if "raw_value" not in updates and "value" not in updates:
+      updates["raw_value"] = _get_val(variable)
     return variable.replace(**updates)
   return variable
 
@@ -234,14 +246,14 @@ def remove_fn(l):
       removed = axis_name in l
       if removed:
         l.remove(axis_name)
-      if len(l) > x.get_value().ndim:
+      if len(l) > _get_val(x).ndim:
         if removed:
           raise ValueError(
-              f"Sharding names {l} still exceed value rank {x.get_value().ndim} after removing scan axis "
+              f"Sharding names {l} still exceed value rank {_get_val(x).ndim} after removing scan axis "
               f"{axis_name!r}; the partition metadata is inconsistent."
           )
         raise ValueError(
-            f"Scan axis {axis_name!r} not found in sharding names {l} for a rank-{x.get_value().ndim} value; "
+            f"Scan axis {axis_name!r} not found in sharding names {l} for a rank-{_get_val(x).ndim} value; "
             "the partition metadata is inconsistent."
         )
       return l
@@ -267,17 +279,17 @@ def _op(x):
     axis_name = x.get_metadata().get(nnx.PARTITION_NAME, name)
     target = x.get_metadata().get("param_scan_axis", pos)
 
-    val = x.get_value()
+    val = _get_val(x)
     if target != 0 and hasattr(val, "ndim") and val.ndim > target:
       x = x.replace(value=jnp.moveaxis(val, 0, target))
 
     def add_fn(l):
       if axis_name not in l:
-        while len(l) < x.get_value().ndim - 1:
+        while len(l) < _get_val(x).ndim - 1:
           l.append(None)
         l.insert(target, axis_name)
       else:
-        while len(l) < x.get_value().ndim:
+        while len(l) < _get_val(x).ndim:
           l.append(None)
       return l
 

src/maxtext/utils/sharding.py

@@ -178,7 +178,12 @@ def remove_size_one_mesh_axis(spec, mesh):
 
 def get_nnx_var_named_sharding_with_scan_axis(v: nnx.Variable, mesh) -> nnx.Variable:
   """Compute NamedSharding for an NNX variable, correctly handling the scan axis."""
-  val = v.get_value()
+  if hasattr(v, "raw_value"):
+    val = v.raw_value
+  elif hasattr(v, "value"):
+    val = v.value
+  else:
+    val = v
   if not hasattr(val, "shape"):
     # `val` is either truly leafless (e.g. optax MaskedNode) or a composite
     # pytree of tensors (e.g. AQT QTensor on serve-mode quantized variables).