[DTensor] Assert DTensorSpec has valid placements

wconstab · wconstab · commit 42d22055a19f · 2025-07-11T12:30:47.000-07:00
This helped identify buggy sharding rules during debugging, why not check it in. ghstack-source-id: 3fde8ba Pull Request resolved: #158133
diff --git a/test/distributed/tensor/test_dtensor_compile.py b/test/distributed/tensor/test_dtensor_compile.py
@@ -298,7 +298,7 @@ def test_dtensor_constructor_w_graph_break(self):
         x = torch.randn(64, 32, requires_grad=True)
         spec = DTensorSpec(
             mesh,
-            (Replicate(), Shard(0)),
+            (Replicate(),),
             tensor_meta=TensorMeta(
                 shape=torch.Size([128, 32]), stride=(32, 1), dtype=x.dtype
             ),
diff --git a/torch/distributed/tensor/_dtensor_spec.py b/torch/distributed/tensor/_dtensor_spec.py
@@ -32,6 +32,10 @@ class DTensorSpec:
     def __post_init__(self) -> None:
         if not isinstance(self.placements, tuple):
             self.placements = tuple(self.placements)
+        if not len(self.placements) == self.mesh.ndim:
+            raise ValueError(
+                f"DTensorSpec requires one placement per mesh dim (mesh.ndim={self.mesh.ndim}), got {self.placements=}"
+            )
         self._hash: Optional[int] = None
 
     def __setattr__(self, attr: str, value: Any) -> None:
