Add sharding propagation support in nnx.eval_shape (clone of #5111)

flax/nnx/transforms/transforms.py

@@ -231,9 +231,21 @@ def to_value_metadata(x):
 
 
 def _to_variable(node):
+  # import here to avoid circular imports
+  from flax.nnx.spmd import get_var_pspec
+
   def to_variable(x):
     if isinstance(x, ValueMetadata):
       var = x.var_type._new(x.value, x.metadata)
+
+      global_mesh = jax.sharding.get_mesh()
+      if global_mesh.axis_sizes == ():
+        global_mesh = None
+      mesh = var.get_metadata("mesh", None) or global_mesh
+      if mesh is not None:
+        pspec = get_var_pspec(var)
+        sharding = jax.sharding.NamedSharding(mesh=mesh, spec=pspec)
+        var.set_value(jax.ShapeDtypeStruct(shape=var.shape, dtype=var.dtype, sharding=sharding))
       return var
     return x
 

tests/nnx/spmd_test.py

@@ -359,6 +359,36 @@ def test_sharding_axis_types(self, mode):
       self.assertEqual(v.sharding.mesh, mesh)
       self.assertEqual(v.sharding.spec, P('row', 'col'))
 
+  def test_eval_shape_with_sharding0(self):
+    # based on https://github.com/google/flax/issues/5110
+    mesh1 = jax.make_mesh((2, 2), ("a", "b"), (jax.sharding.AxisType.Auto, jax.sharding.AxisType.Auto))
+    mesh2 = jax.make_mesh((1, 4), ("c", "d"), (jax.sharding.AxisType.Auto, jax.sharding.AxisType.Auto))
+
+    class Model(nnx.Module):
+        def __init__(self):
+            self.p1 = nnx.Linear(16, 16, rngs=nnx.Rngs(0), kernel_metadata={"out_sharding": ("a", "b"), "mesh": mesh1})
+            self.p2 = nnx.Linear(16, 16, rngs=nnx.Rngs(0), kernel_metadata={"out_sharding": ("c", "d"), "mesh": mesh2})
+
+    abs_model = nnx.eval_shape(lambda: Model())
+    assert isinstance(abs_model.p1.kernel.sharding, jax.sharding.NamedSharding)
+    assert abs_model.p1.kernel.sharding.mesh is mesh1
+    assert abs_model.p1.kernel.sharding.spec == jax.P("a", "b")
+    assert isinstance(abs_model.p2.kernel.sharding, jax.sharding.NamedSharding)
+    assert abs_model.p2.kernel.sharding.mesh is mesh2
+    assert abs_model.p2.kernel.sharding.spec == jax.P("c", "d")
+
+  def test_eval_shape_with_sharding1(self):
+    class Model(nnx.Module):
+        def __init__(self):
+            self.linear = nnx.Linear(10, 10, rngs=nnx.Rngs(0), kernel_metadata={"out_sharding": ("a", "b")})
+
+    mesh = jax.make_mesh((2, 2), ("a", "b"), (jax.sharding.AxisType.Auto, jax.sharding.AxisType.Auto))
+    with jax.set_mesh(mesh):
+        abs_model = nnx.eval_shape(lambda: Model())
+    assert isinstance(abs_model.linear.kernel.sharding, jax.sharding.NamedSharding)
+    assert abs_model.linear.kernel.sharding.mesh is mesh
+    assert abs_model.linear.kernel.sharding.spec == jax.P("a", "b")
+
 def has_sharding_spec(array):
     sharding = array.sharding
     if hasattr(sharding, 'spec'):