Add the axis_types as Auto for the test_qkvo_partition_spec unit test

Author: judyzhaoxinwu

axlearn/common/attention_test.py

@@ -2651,7 +2651,15 @@ def test_qkvo_partition_spec(self, kv_part):
             self.skipTest(f"Unsupported mesh shape {mesh_shape}")
         model_dim = 16
         num_heads = 4
-        mesh = jax.make_mesh(mesh_shape, axis_names=("fsdp", "seq", "model"))
+        mesh = jax.make_mesh(
+            mesh_shape,
+            axis_names=("fsdp", "seq", "model"),
+            axis_types=(
+                jax.sharding.AxisType.Auto,
+                jax.sharding.AxisType.Auto,
+                jax.sharding.AxisType.Auto,
+            ),
+        )
         q_part = ("fsdp", "seq", "model", None)
         o_part = ("fsdp", "seq", None)
 

axlearn/common/utils.py

@@ -637,14 +637,7 @@ def with_sharding_constraint(x: Tensor, shardings):
     mesh = thread_resources.env.physical_mesh
     if mesh.empty or mesh.size == 1:
         return x
-    try:
-        return jax.lax.with_sharding_constraint(x, shardings)
-    except ValueError as e:
-        logging.warning("with_sharding_constraint failed: %s. Try to re-shard.", e)
-        if "reshard" in str(e) or "Explicit" in str(e):
-            named_sharding = jax.sharding.NamedSharding(mesh, shardings)
-            return jax.lax.reshard(x, named_sharding)
-        raise e
+    return jax.lax.with_sharding_constraint(x, shardings)
 
 
 def maybe_shard(x: NestedTensor, partition_spec: Optional[PartitionSpec]) -> NestedTensor: