From a0ce6ec4908b8a8515a3b3fe78b3945c43a3d3b0 Mon Sep 17 00:00:00 2001
From: Sohum Trivedi <trivsohum@gmail.com>
Date: Fri, 17 Jul 2026 10:19:43 -0700
Subject: [PATCH] Fix nnx.Dropout crash when broadcast_dims includes a sharded
 axis

Dropout passes the input's full sharding as out_sharding to
random.bernoulli, but the mask shape has broadcast_dims collapsed to 1.
Under JAX explicit sharding a size-1 axis cannot be partitioned across
devices, so Dropout with broadcast_dims on a sharded axis raised
'Sharding spec ... does not evenly divide the dimension size 1'.

Derive the mask sharding from the input sharding with the broadcast
dims left unpartitioned, and reshard the broadcasted mask back to the
input sharding before lax.select. Behavior is unchanged when
broadcast_dims is empty or the broadcast axes are unsharded.
---
 flax/nnx/nn/stochastic.py | 15 ++++++++++++++-
 tests/nnx/spmd_test.py    | 21 +++++++++++++++++++++
 2 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/flax/nnx/nn/stochastic.py b/flax/nnx/nn/stochastic.py
index 6d03e7353..10de82bd8 100644
--- a/flax/nnx/nn/stochastic.py
+++ b/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(
diff --git a/tests/nnx/spmd_test.py b/tests/nnx/spmd_test.py
index 5f075ce77..6d3022883 100644
--- a/tests/nnx/spmd_test.py
+++ b/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):