flax

Fix nnx.Dropout crash when broadcast_dims includes a sharded axis

#5528

Open