Fix DTensor handling of conjugate bit.

ezyang · ezyang · commit 89cd4c9b197e · 2025-07-10T08:24:25.000-07:00
Fixes #130646 specifically for DTensor Fixes pytorch/torchtitan#267 Signed-off-by: Edward Z. Yang <ezyang@meta.com> ghstack-source-id: 5b5124e Pull-Request: #158030
diff --git a/torch/distributed/tensor/_api.py b/torch/distributed/tensor/_api.py
@@ -269,6 +269,13 @@ def __new__(
         # new method instruct wrapper tensor from local_tensor and add
         # placement spec, it does not do actual distribution
         assert spec.tensor_meta is not None, "TensorMeta should not be None!"
+        extra_dispatch_keys = torch._C.DispatchKeySet.from_raw_repr(0)
+        if torch._C._dispatch_keys(local_tensor).has(torch._C.DispatchKey.Conjugate):
+            extra_dispatch_keys = extra_dispatch_keys.add(
+                torch._C.DispatchKey.Conjugate
+            )
+        if torch._C._dispatch_keys(local_tensor).has(torch._C.DispatchKey.Negative):
+            extra_dispatch_keys = extra_dispatch_keys.add(torch._C.DispatchKey.Negative)
         r = torch.Tensor._make_wrapper_subclass(
             cls,
             spec.tensor_meta.shape,
@@ -277,6 +284,7 @@ def __new__(
             device=local_tensor.device,
             layout=local_tensor.layout,
             requires_grad=requires_grad,
+            _extra_dispatch_keys=extra_dispatch_keys,
         )
 
         r._spec = spec
