reshard back for 3d

ClassicLarry · ClassicLarry · commit 202715df5b6e · 2026-03-04T15:56:09.000-08:00
diff --git a/lib/levanter/src/levanter/optim/grugmuon.py b/lib/levanter/src/levanter/optim/grugmuon.py
@@ -135,6 +135,7 @@ def transform_array(x):
                 updated = jax.vmap(
                     lambda m: _newtonschulz_core(m, steps=steps, eps=muon_eps, coefficient_type=coefficient_type)
                 )(x)
+                updated = reshard(updated, current_spec)
                 # Layout per slice is (fan_in, fan_out)
                 fan_in, fan_out = updated.shape[1], updated.shape[2]
             else:
