Track explicit sharding in Variables

Author: samanklesaria

flax/nnx/spmd.py

@@ -74,10 +74,10 @@ def remove_field(fields, index, value):
 
   def _remove_axis(x: tp.Any):
     if isinstance(x, variablelib.Variable):
-      if hasattr(x, 'out_sharding') and x.out_sharding is not None:
-        x.set_metadata(
-          out_sharding=remove_field(x.out_sharding, index, axis_name)
-        )
+      if hasattr(x, 'out_sharding') and x.out_sharding:
+          x.set_metadata(
+            out_sharding=remove_field(x.out_sharding, index, axis_name)
+          )
 
       for k, v in other_meta.items():
         if hasattr(x, k) and (t := getattr(x, k)) and isinstance(t, tuple):
@@ -174,4 +174,4 @@ def get_abstract_model(init_fn, mesh):
       lambda a, s: jax.ShapeDtypeStruct(a.shape, a.dtype, sharding=s),
       abs_state, get_named_sharding(abs_state, mesh)
     )
-  return gdef, abs_state
+  return gdef, abs_state

flax/nnx/variablelib.py

@@ -28,6 +28,7 @@
 from flax.core import spmd as core_spmd
 from flax.nnx import reprlib, tracers, visualization
 from flax.typing import MISSING, Missing, SizeBytes
+from jax.sharding import NamedSharding
 import jax
 from jax._src.state.types import AbstractRef
 import jax.experimental
@@ -1418,6 +1419,10 @@ def __init__(
     if ref:
       value = jax.new_ref(value)  # type: ignore
     object.__setattr__(self, '_raw_value', value)
+    if ((not hasattr(metadata, 'out_sharding')) and hasattr(value, 'sharding') and
+      isinstance(value.sharding, NamedSharding) and any(value.sharding.spec)):
+      metadata['out_sharding'] = value.sharding.spec
+      metadata['mesh'] = value.sharding.mesh
 
   @property
   def _can_update(self) -> bool: