Added test for catching thrown error in spec

Author: aws-cph

test/spmd/test_xla_dtensor_spec_conversion.py

@@ -212,6 +212,23 @@ def test_spec_invalidation_on_resharding(self):
     assert resharded_tensor._spec is not initial_spec
     assert resharded_tensor._spec.placements[1].dim == 1
 
+  def test_auto_wrapped_tensor_spec_failure(self):
+    """Test that auto-wrapped tensors fail when accessing _spec property.
+    
+    Auto-wrapped tensors are created through operations that trigger __torch_dispatch__
+    but don't yet have access to the sharding propagation done through open xla,
+    causing ._spec to fail. 
+    """
+    device_count = xr.global_runtime_device_count()
+    mesh = DeviceMesh("xla", torch.arange(device_count))
+    tensor = torch.randn(4, 4)
+    sharded_tensor = distribute_tensor(tensor, mesh, [Shard(0)])
+
+    auto_wrapped = sharded_tensor + sharded_tensor
+
+    with self.assertRaises(ValueError):
+      _ = auto_wrapped._spec
+
 
 if __name__ == '__main__':
   test = unittest.main()

torch_xla/distributed/spmd/xla_sharded_tensor.py

@@ -205,7 +205,11 @@ def _spec(self):
       mesh = DeviceMesh("xla",
                         torch.tensor(device_list).reshape(self.mesh_shape))
     else:
-      raise ValueError("mesh_shape must be specified to create DTensorSpec")
+      raise ValueError(
+          "mesh_shape must be specified to create DTensorSpec. "
+          "If this tensor was created through torch operations, it may be auto-wrapped. "
+          "Use wrap_as_sharded_tensor() to set mesh_shape before accessing _spec. "
+      )
 
     # use existing partition_spec
     if self.partition_spec is not None:
@@ -220,7 +224,11 @@ def _spec(self):
         placements.append(
             Shard(tensor_dim) if tensor_dim is not None else Replicate())
     else:
-      raise ValueError("partition_spec must be specified to create DTensorSpec")
+      raise ValueError(
+          "partition_spec must be specified to create DTensorSpec. "
+          "If this tensor was created through torch operations, it may be auto-wrapped. "
+          "Use wrap_as_sharded_tensor() to set partition_spec before accessing _spec. "
+      )
 
     # tensor metadata
     tensor_meta = TensorMeta(