Add a bunch of ops trying to get gradients to work.

Author: hameerabbasi

src/complex_tensor/ops/aten.py

@@ -81,6 +81,8 @@ def is_pinned_impl(self: ComplexTensor, device: torch.device | None = None) -> b
 detach_impl = register_simple(aten.detach)
 select_impl = register_simple(aten.select)
 squeeze_impl = register_simple(aten.squeeze)
+zero__impl = register_simple(aten.zero_)
+transpose_impl = register_simple(aten.transpose)
 
 # TODO (hameerabbasi): Not being tested
 copy_impl = register_force_test(aten.copy, register_simple(aten.copy))
@@ -533,14 +535,14 @@ def copy__impl(self: ComplexTensor, src, *args, **kwargs):
 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)
+    self_re, self_im = split_complex_tensor(self)
+    if dtype is not None and dtype not in COMPLEX_TO_REAL:
+        return self_re.new_zeros(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)
+        dtype = COMPLEX_TO_REAL[dtype]
+    re = self_re.new_zeros(size, dtype=dtype, **kwargs)
+    im = self_im.new_zeros(size, dtype=dtype, **kwargs)
 
     return ComplexTensor(re, im)
 
@@ -561,7 +563,7 @@ def allclose_impl(
     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))
+    return torch.all(torch.isclose(input, other, rtol=rtol, atol=atol, equal_nan=equal_nan)).item()
 
 
 @register_complex(aten.stack)
@@ -570,3 +572,17 @@ def stack_impl(self: list[ComplexTensor], *args, **kwargs) -> ComplexTensor:
     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)
+
+
+@register_complex(aten.randn_like)
+def randn_like_impl(self: ComplexTensor, *, dtype=None, **kwargs) -> ComplexTensor | torch.Tensor:
+    if dtype is not None and dtype not in COMPLEX_TO_REAL:
+        return torch.randn_like(self.re, dtype=dtype, **kwargs)
+
+    if dtype is not None:
+        dtype = COMPLEX_TO_REAL[dtype]
+
+    self_re, self_im = split_complex_tensor(self)
+    ret_re = torch.randn_like(self_re, dtype=dtype, **kwargs) / 2
+    ret_im = torch.randn_like(self_im, dtype=dtype, **kwargs) / 2
+    return ComplexTensor(ret_re, ret_im)

src/complex_tensor/ops/prims.py

@@ -3,6 +3,7 @@
 from ..complex_tensor import ComplexTensor
 from ._common import (
     complex_to_real_dtype,
+    register_complex,
     register_force_test,
     split_complex_tensor,
 )
@@ -18,3 +19,10 @@ def convert_element_type_impl(x: ComplexTensor, dtype: torch.dtype) -> ComplexTe
     v_out = prims.convert_element_type(v, dtype)
 
     return ComplexTensor(u_out, v_out)
+
+
+@register_complex(prims.conj_physical)
+@register_complex(prims.conj)
+def conj_physical_impl(self: ComplexTensor) -> ComplexTensor:
+    re, im = split_complex_tensor(self)
+    return ComplexTensor(re, -im)

src/complex_tensor/test/test_ops.py

@@ -118,6 +118,7 @@ def test_fn_grad(self, device, dtype, op: OpInfo) -> None:
             self.skipTest("Skipped! Dtype is not in supported backward dtypes!")
         else:
             with ComplexDispatchMode():
+                op.gradcheck_fast_mode = False
                 self._grad_test_helper(device, dtype, op, op.get_op())