Add gradient checks: first attempt.

Author: hameerabbasi

src/complex_tensor/ops/_common.py

@@ -4,7 +4,8 @@
 import torch
 from torch._ops import OpOverloadPacket
 from torch._refs import is_complex
-from torch.utils._pytree import tree_flatten, tree_unflatten
+from torch.utils._python_dispatch import TorchDispatchMode
+from torch.utils._pytree import tree_flatten, tree_map, tree_unflatten
 
 from ..complex_tensor import ComplexTensor
 
@@ -131,7 +132,7 @@ def ordered_impl(*args, **kwargs):
 
 def register_binary_nonlinear(op: OpType) -> Callable:
     def impl(lhs: ComplexTensor, rhs: ComplexTensor, *args, **kwargs) -> ComplexTensor:
-        a_r, a_i = split_complex_tensor(lhs)
+        a_r, a_i = split_complex_arg(lhs)
         b_r, b_i = split_complex_arg(rhs)
         out_dt, (a_r, a_i, b_r, b_i) = promote_real_cpu_tensors(a_r, a_i, b_r, b_i)
         real = op(a_r, b_r, *args, **kwargs) - op(a_i, b_i, *args, **kwargs)
@@ -161,3 +162,30 @@ def impl(self: ComplexTensor, *args, **kwargs) -> ComplexTensor:
     impl.__qualname__ = func_name
 
     return register_complex(op, impl)
+
+
+def _as_complex_tensor(arg: torch.Tensor | Any) -> torch.Tensor | ComplexTensor | Any:
+    if (
+        not isinstance(arg, ComplexTensor)
+        and isinstance(arg, torch.Tensor)
+        and arg.dtype in COMPLEX_TO_REAL
+    ):
+        return ComplexTensor.from_interleaved(arg)
+    return arg
+
+
+def _as_interleaved(arg: ComplexTensor | Any) -> torch.Tensor | Any:
+    if isinstance(arg, ComplexTensor):
+        return arg.as_interleaved()
+    return arg
+
+
+class ComplexDispatchMode(TorchDispatchMode):
+    def __torch_dispatch__(self, func, types, args=(), kwargs=None):
+        if kwargs is None:
+            kwargs = {}
+
+        args = tree_map(_as_complex_tensor, args)
+        kwargs = tree_map(_as_complex_tensor, kwargs)
+
+        return tree_map(_as_interleaved, func(*args, **kwargs))

src/complex_tensor/ops/aten.py

@@ -33,7 +33,7 @@ def impl(lhs: ComplexTensor, rhs: ComplexTensor, *args, **kwargs) -> ComplexTens
         alpha = kwargs.pop("alpha", None)
         if alpha is not None:
             return impl_with_alpha(lhs, rhs, *args, alpha=alpha, **kwargs)
-        a_r, a_i = split_complex_tensor(lhs)
+        a_r, a_i = split_complex_arg(lhs)
         b_r, b_i = split_complex_arg(rhs)
         out_dt, (a_r, a_i, b_r, b_i) = promote_real_cpu_tensors(a_r, a_i, b_r, b_i)
         u = op(a_r, b_r, *args, **kwargs)
@@ -78,6 +78,9 @@ def is_pinned_impl(self: ComplexTensor, device: torch.device | None = None) -> b
 index_select_impl = register_simple(aten.index_select)
 split_with_sizes_impl = register_simple(aten.split_with_sizes)
 cumsum_impl = register_simple(aten.cumsum)
+detach_impl = register_simple(aten.detach)
+select_impl = register_simple(aten.select)
+squeeze_impl = register_simple(aten.squeeze)
 
 # TODO (hameerabbasi): Not being tested
 copy_impl = register_force_test(aten.copy, register_simple(aten.copy))
@@ -502,3 +505,68 @@ def nonzero_impl(self: ComplexTensor, other: ComplexTensor, *args, **kwargs) ->
 @register_complex(aten.logical_not)
 def logical_not_impl(self: ComplexTensor, *args, **kwargs) -> torch.Tensor:
     return torch.logical_not(elemwise_nonzero(self), *args, **kwargs)
+
+
+@register_complex(aten.view_as_real)
+def view_as_real_impl(self: ComplexTensor) -> torch.Tensor:
+    re, im = split_complex_tensor(self)
+    return torch.stack([re, im], dim=-1)
+
+
+@register_complex(aten.linalg_vector_norm)
+def linalg_vector_norm_impl(self: ComplexTensor, *args, **kwargs) -> torch.Tensor:
+    return torch.linalg.vector_norm(torch.abs(self), *args, **kwargs)
+
+
+@register_force_test(aten.copy_)
+def copy__impl(self: ComplexTensor, src, *args, **kwargs):
+    self_re, self_im = split_complex_tensor(self)
+    src_re, src_im = split_complex_arg(src)
+
+    ret_re = self_re.copy_(src_re, *args, **kwargs)
+    ret_im = self_im.copy_(src_im, *args, **kwargs)
+
+    return ComplexTensor(ret_re, ret_im)
+
+
+@register_complex(aten.new_zeros)
+def new_zeros_impl(
+    self: ComplexTensor, size, *, dtype=None, **kwargs
+) -> ComplexTensor | torch.Tensor:
+    if dtype is not None and torch.dtype(dtype) not in COMPLEX_TO_REAL:
+        return self.re.new_zeros(self, size, dtype=dtype, **kwargs)
+
+    if dtype is not None:
+        dtype = COMPLEX_TO_REAL[torch.dtype(dtype)]
+
+    re = self.re.new_zeros(size, dtype=dtype, **kwargs)
+    im = self.im.new_zeros(size, dtype=dtype, **kwargs)
+
+    return ComplexTensor(re, im)
+
+
+@register_complex(aten._local_scalar_dense)
+def _local_scalar_dense_impl(self: ComplexTensor, *args, **kwargs) -> complex:
+    x, y = split_complex_tensor(self)
+    u = aten._local_scalar_dense(x, *args, **kwargs)
+    v = aten._local_scalar_dense(y, *args, **kwargs)
+    return complex(u, v)
+
+
+@register_complex(aten.allclose)
+def allclose_impl(
+    input: torch.Tensor,
+    other: torch.Tensor,
+    rtol: float = 1e-05,
+    atol: float = 1e-08,
+    equal_nan: bool = False,
+) -> torch.Tensor:
+    return torch.all(torch.isclose(input, other, rtol=rtol, atol=atol, equal_nan=equal_nan))
+
+
+@register_complex(aten.stack)
+def stack_impl(self: list[ComplexTensor], *args, **kwargs) -> ComplexTensor:
+    re_im_tuples = [split_complex_arg(self_i) for self_i in self]
+    u = torch.stack([c[0] for c in re_im_tuples], *args, **kwargs)
+    v = torch.stack([c[1] for c in re_im_tuples], *args, **kwargs)
+    return ComplexTensor(u, v)

src/complex_tensor/test/test_ops.py

@@ -4,11 +4,22 @@
 from torch._ops import OpOverload
 from torch.testing._internal.common_device_type import instantiate_device_type_tests, ops
 from torch.testing._internal.common_methods_invocations import op_db
-from torch.testing._internal.common_utils import parametrize, run_tests
+from torch.testing._internal.common_utils import (
+    TestGradients,
+    parametrize,
+    run_tests,
+    unMarkDynamoStrictTest,
+)
 from torch.testing._internal.opinfo.core import OpInfo
 
 from complex_tensor.ops import COMPLEX_OPS_TABLE, FORCE_TEST_LIST
-from complex_tensor.test.utils import COMPLEX_DTYPES, TestCase, TestDescriptor, _as_complex_tensor
+from complex_tensor.test.utils import (
+    COMPLEX_DTYPES,
+    ComplexDispatchMode,
+    TestCase,
+    TestDescriptor,
+    _as_complex_tensor,
+)
 
 torch._dynamo.config.recompile_limit = float("inf")
 torch._dynamo.config.accumulated_recompile_limit = float("inf")
@@ -99,7 +110,19 @@ def actual(subclass_sample=subclass_sample):
             self.assertSameResult(expected, actual, ignore_exc_types=compile)
 
 
+@unMarkDynamoStrictTest
+class TestComplexBwdGradients(TestGradients):
+    @ops(implemented_op_db, allowed_dtypes=list(COMPLEX_DTYPES))
+    def test_fn_grad(self, device, dtype, op: OpInfo) -> None:
+        if dtype not in op.supported_backward_dtypes(torch.device(device).type):
+            self.skipTest("Skipped! Dtype is not in supported backward dtypes!")
+        else:
+            with ComplexDispatchMode():
+                self._grad_test_helper(device, dtype, op, op.get_op())
+
+
 instantiate_device_type_tests(TestComplexTensor, globals())
+instantiate_device_type_tests(TestComplexBwdGradients, globals())
 
 if __name__ == "__main__":
     run_tests()

src/complex_tensor/test/utils.py

@@ -8,9 +8,9 @@
 from torch.testing._internal.common_utils import TestCase as PytorchTestCase
 from torch.utils._pytree import tree_flatten
 
-from complex_tensor.complex_tensor import ComplexTensor
+from complex_tensor.ops._common import COMPLEX_TO_REAL, _as_complex_tensor
 
-COMPLEX_DTYPES = {torch.complex128, torch.complex64, torch.complex32}
+COMPLEX_DTYPES = set(COMPLEX_TO_REAL)
 
 
 @dataclass(frozen=True)
@@ -35,16 +35,6 @@ def matches(self, other: TestDescriptor) -> bool:
         return True
 
 
-def _as_complex_tensor(arg):
-    if (
-        not isinstance(arg, ComplexTensor)
-        and isinstance(arg, torch.Tensor)
-        and arg.dtype in COMPLEX_DTYPES
-    ):
-        return ComplexTensor.from_interleaved(arg)
-    return arg
-
-
 class TestCase(PytorchTestCase):
     def assertSameResult(
         self,