Fix nnx.Dropout crash when broadcast_dims includes a sharded axis

flax/nnx/nn/stochastic.py

@@ -159,10 +159,23 @@ def __call__(
     broadcast_shape = list(inputs.shape)
     for dim in self.broadcast_dims:
       broadcast_shape[dim] = 1
+    inputs_sharding = jax.typeof(inputs).sharding
+    # The mask has size 1 along the broadcast dims, so it cannot be
+    # partitioned there; drop those axes from the mask's sharding.
+    mask_spec = list(inputs_sharding.spec) + [None] * (
+      inputs.ndim - len(inputs_sharding.spec)
+    )
+    for dim in self.broadcast_dims:
+      mask_spec[dim] = None
+    mask_sharding = inputs_sharding.update(
+      spec=jax.sharding.PartitionSpec(*mask_spec)
+    )
     mask = random.bernoulli(
-      key, p=keep_prob, shape=broadcast_shape, out_sharding=jax.typeof(inputs).sharding
+      key, p=keep_prob, shape=broadcast_shape, out_sharding=mask_sharding
     )
     mask = jnp.broadcast_to(mask, inputs.shape)
+    if mask_sharding != inputs_sharding:
+      mask = jax.sharding.reshard(mask, inputs_sharding)
     return lax.select(mask, inputs / keep_prob, jnp.zeros_like(inputs))
 
   def set_view(

tests/nnx/spmd_test.py

@@ -353,6 +353,27 @@ def func(x, rngs):
 
         assert 'float32[2@X,4]' in str(jax.typeof(func(sharded_array, nnx.Rngs(0))))
 
+  def test_out_sharding_dropout_broadcast_sharded_axis(self):
+    mesh = jax.make_mesh((2, 2), ("X", "Y"), axis_types=(AxisType.Explicit, AxisType.Explicit))
+    with jax.set_mesh(mesh):
+      replicated_array = (jnp.arange(8) + 1).reshape(2, 4).astype(jnp.float32)
+      sharded_array = reshard(replicated_array, P("X", None))
+      # broadcast_dims includes the sharded axis 0
+      layer = nnx.Dropout(
+        rate=0.5, deterministic=False, broadcast_dims=(0,), rngs=nnx.Rngs(0)
+      )
+      out = layer(sharded_array)
+      assert 'float32[2@X,4]' in str(jax.typeof(out))
+      # the dropout mask must be shared along the broadcast axis
+      out_host = jax.device_get(out)
+      assert ((out_host[0] == 0.0) == (out_host[1] == 0.0)).all()
+
+      @jax.jit
+      def func(x, rngs):
+        return layer(x, rngs=rngs)
+
+      assert 'float32[2@X,4]' in str(jax.typeof(func(sharded_array, nnx.Rngs(0))))
+
   def test_out_sharding_mha(self):
     mesh = jax.make_mesh((2, 2), ("fsdp", "tp"), axis_types=(AxisType.Explicit, AxisType.Explicit))
     with jax.set_mesh(mesh):