Support conj bit, neg bit

Author: zou3519

functorch/csrc/BatchedTensorImpl.cpp

@@ -24,6 +24,8 @@ BatchedTensorImpl::BatchedTensorImpl(Tensor value, BatchDims bdims)
   , value_(std::move(value))
   , bdims_(std::move(bdims))
 {
+  // TODO: I don't think this ctor gets used.
+  TORCH_INTERNAL_ASSERT(false);
   TORCH_INTERNAL_ASSERT(value_.defined());
   set_storage_access_should_throw();
   set_has_contiguity_policy(HasContiguityPolicy::CustomBehavior);
@@ -147,10 +149,7 @@ const char* BatchedTensorImpl::tensorimpl_type_name() const {
 }
 
 Tensor makeBatched(const Tensor& tensor, BatchDims bdims) {
-  DispatchKeySet key_set;
-  if (tensor.is_cuda()) {
-    key_set = key_set.add(DispatchKey::CUDA);
-  }
+  DispatchKeySet key_set = getKeysToPropagateToWrapper(tensor);
   auto* batched = maybeGetBatchedImpl(tensor);
   if (batched) {
     auto requested_level = bdims.back().level();

functorch/csrc/BatchedTensorImpl.h

@@ -160,5 +160,18 @@ TORCH_API Tensor addBatchDim(const Tensor& tensor, int64_t level, int64_t dim);
 // See NOTE: [vmap-incompatible in-place operations] for the definition of this.
 TORCH_API bool inplaceIsVmapCompatible(const Tensor& self, const Tensor& other);
 
+constexpr DispatchKeySet kKeysToPropagateToWrapper({
+  DispatchKey::Negative,
+  DispatchKey::Conjugate,
+  DispatchKey::XLA,
+  DispatchKey::CUDA,
+  DispatchKey::CPU,
+});
+
+inline DispatchKeySet getKeysToPropagateToWrapper(const Tensor& tensor, DispatchKeySet to_propagate=kKeysToPropagateToWrapper) {
+  auto key_set = tensor.unsafeGetTensorImpl()->key_set();
+  return key_set & kKeysToPropagateToWrapper;
+}
+
 }
 }

functorch/csrc/TensorWrapper.cpp

@@ -62,15 +62,9 @@ void dumpTensorCout(const Tensor& tensor) {
 }
 
 c10::intrusive_ptr<TensorWrapper> makeTensorWrapperPtr(const Tensor& tensor, int64_t level, bool should_be_alive) {
-  // TODO: denylist non-cuda/cpu backends to avoid funny business
-  DispatchKeySet key_set;
-  if (tensor.is_cuda()) {
-    key_set = key_set.add(DispatchKey::CUDA);
-    key_set = key_set.add(DispatchKey::AutogradCUDA);
-  } else {
-    key_set = key_set.add(DispatchKey::CPU);
-    key_set = key_set.add(DispatchKey::AutogradCPU);
-  }
+  auto keys_to_propagate = kKeysToPropagateToWrapper | DispatchKeySet({
+      DispatchKey::AutogradCPU, DispatchKey::AutogradCUDA, DispatchKey::AutogradXLA});
+  auto key_set = getKeysToPropagateToWrapper(tensor, keys_to_propagate);
   key_set = key_set.add(kGradWrapperKey);
   if (should_be_alive) {
     return c10::make_intrusive<TensorWrapper>(key_set, tensor, level, getLifeHandleForLevel(level));
@@ -85,15 +79,9 @@ Tensor makeTensorWrapper(const Tensor& tensor, int64_t level) {
     TORCH_INTERNAL_ASSERT(wrapped->level() < level);
   }
 
-  // TODO: denylist non-cuda/cpu backends to avoid funny business
-  DispatchKeySet key_set;
-  if (tensor.is_cuda()) {
-    key_set = key_set.add(DispatchKey::CUDA);
-    key_set = key_set.add(DispatchKey::AutogradCUDA);
-  } else {
-    key_set = key_set.add(DispatchKey::CPU);
-    key_set = key_set.add(DispatchKey::AutogradCPU);
-  }
+  auto keys_to_propagate = kKeysToPropagateToWrapper | DispatchKeySet({
+      DispatchKey::AutogradCPU, DispatchKey::AutogradCUDA, DispatchKey::AutogradXLA});
+  auto key_set = getKeysToPropagateToWrapper(tensor, keys_to_propagate);
   key_set = key_set.add(kGradWrapperKey);
   auto life_handle = getLifeHandleForLevel(level);
   auto result = at::detail::make_tensor<TensorWrapper>(key_set, tensor, level, std::move(life_handle));

test/test_eager_transforms.py

@@ -200,6 +200,16 @@ def f(x):
         result, vjp_fn = vjp(f, torch.tensor(1.))
         vjp_fn(result)
 
+    def test_conj_bit(self):
+        x = torch.tensor(1+1j)
+        def foo(x):
+            assert not x.is_conj()
+            y = x.conj()
+            assert y.is_conj()
+            return y
+        res = grad(foo)(x)
+        self.assertEqual(res, torch.ones_like(res))
+
     def test_composed_with_autograd(self, device):
         x = torch.randn([], requires_grad=True, device=device)
 

test/test_vmap.py

@@ -2618,6 +2618,16 @@ def test_one_hot(self):
             for loop_out, batched_out in get_fallback_and_vmap_exhaustive(F.one_hot, args, {}):
                 self.assertEqual(loop_out, batched_out)
 
+    def test_conj_bit(self):
+        x = torch.tensor([1+1j, 2+1j])
+        def foo(x):
+            assert not x.is_conj()
+            y = x.conj()
+            assert y.is_conj()
+            return y
+        res = vmap(foo)(x)
+        self.assertEqual(res, x.conj())
+
     def test_mode_key(self):
         def vmap_f(x):
             return x + torch.randn(())