Update on "[dtensor] Introduce full_tensor API to DTensor"

wanchaol · wanchaol · commit 1cdc18758009 · 2023-10-27T14:42:47.000-07:00
This PR introduces a `full_tensor` API to DTensor, there were so many
callsites that exercises the `redistribute(replicate)` path and I feel
it deserves a separate API, mostly just a syntactic sugar

[ghstack-poisoned]
diff --git a/test/distributed/_tensor/test_random_ops.py b/test/distributed/_tensor/test_random_ops.py
@@ -125,7 +125,9 @@ def test_deterministic_dropout_1d(self):
         dtensor = dropout(dtensor)
 
         # allgather the local tensors
-        local_tensor = dtensor.full_tensor()
+        local_tensor = funcol.all_gather_tensor(
+            dtensor.to_local(), gather_dim=0, group=(device_mesh, 0)
+        )
 
         # compare with local tensors from other ranks
         self_slice = slice(4 * self.rank, 4 * self.rank + 4)
