[DTensor] Used new placements for neg dim in from_local

Andrew Gu · Andrew Gu · commit 43a9a1fe209e · 2023-11-20T08:47:15.000-08:00
[ghstack-poisoned]
diff --git a/torch/distributed/_tensor/api.py b/torch/distributed/_tensor/api.py
@@ -343,6 +343,14 @@ def from_local(
         # set default placements to replicated if not specified
         if placements is None:
             placements = [Replicate() for _ in range(device_mesh.ndim)]
+        else:
+            placements = list(placements)
+            for idx, placement in enumerate(placements):
+                # normalize shard dim to be positive
+                if placement.is_shard():
+                    placement = cast(Shard, placement)
+                    if placement.dim < 0:
+                        placements[idx] = Shard(placement.dim + local_tensor.ndim)
 
         # `from_local` is differentiable, and the gradient of the dist tensor this function
         # created should flow back the gradients to the local_tensor, so we call an autograd
