Merge pull request #7 from Emrys365/fixup_wangyou

Fix up bugs/interfaces in matmul and squeeze, and add .unbind to ComplexTensor

torch_complex/functional.py

@@ -218,6 +218,7 @@ def matmul(
     if isinstance(a, ComplexTensor) and isinstance(b, ComplexTensor):
         return a @ b
     elif not isinstance(a, ComplexTensor) and isinstance(b, ComplexTensor):
+        o_real = torch.matmul(a, b.real)
         o_imag = torch.matmul(a, b.imag)
     elif isinstance(a, ComplexTensor) and not isinstance(b, ComplexTensor):
         return a @ b

torch_complex/tensor.py

@@ -619,8 +619,11 @@ def ndim(self):
     def sqrt(self) -> "ComplexTensor":
         return self ** 0.5
 
-    def squeeze(self, dim) -> "ComplexTensor":
-        return ComplexTensor(self.real.squeeze(dim), self.imag.squeeze(dim))
+    def squeeze(self, dim=None) -> "ComplexTensor":
+        if dim is None:
+            return ComplexTensor(self.real.squeeze(), self.imag.squeeze())
+        else:
+            return ComplexTensor(self.real.squeeze(dim), self.imag.squeeze(dim))
 
     def sum(self, *args, **kwargs) -> "ComplexTensor":
         """
@@ -668,6 +671,14 @@ def type(self, *args, **kwargs) -> str:
                 self.real.type(*args, **kwargs), self.imag.type(*args, **kwargs)
             )
 
+    def unbind(self, dim=0) -> "ComplexTensor":
+        return tuple(
+            map(
+                lambda x: ComplexTensor(*x),
+                zip(self.real.unbind(dim=dim), self.imag.unbind(dim=dim))
+            )
+        )
+
     def unfold(self, dim, size, step):
         return ComplexTensor(
             self.real.unfold(dim, size, step), self.imag.unfold(dim, size, step)