Merge pull request #11 from yoyololicon/feat/remove-einsum

feat: remove einsum for efficiency

Author: fkodom

fft_conv_pytorch/fft_conv.py

@@ -13,16 +13,17 @@ def complex_matmul(a: Tensor, b: Tensor, groups: int = 1) -> Tensor:
     # dimensions. Dimensions 3 and higher will have the same shape after multiplication.
     # We also allow for "grouped" multiplications, where multiple sections of channels
     # are multiplied independently of one another (required for group convolutions).
-    scalar_matmul = partial(torch.einsum, "agc..., gbc... -> agb...")
     a = a.view(a.size(0), groups, -1, *a.shape[2:])
     b = b.view(groups, -1, *b.shape[1:])
 
-    # Compute the real and imaginary parts independently, then manually insert them
-    # into the output Tensor.  This is fairly hacky but necessary for PyTorch 1.7.0,
-    # because Autograd is not enabled for complex matrix operations yet.  Not exactly
-    # idiomatic PyTorch code, but it should work for all future versions (>= 1.7.0).
-    real = scalar_matmul(a.real, b.real) - scalar_matmul(a.imag, b.imag)
-    imag = scalar_matmul(a.imag, b.real) + scalar_matmul(a.real, b.imag)
+    a = torch.movedim(a, 2, a.dim() - 1).unsqueeze(-2)
+    b = torch.movedim(b, (1, 2), (b.dim() - 1, b.dim() - 2))
+
+    # complex value matrix multiplication
+    real = a.real @ b.real - a.imag @ b.imag
+    imag = a.imag @ b.real + a.real @ b.imag
+    real = torch.movedim(real, real.dim() - 1, 2).squeeze(-1)
+    imag = torch.movedim(imag, imag.dim() - 1, 2).squeeze(-1)
     c = torch.zeros(real.shape, dtype=torch.complex64, device=a.device)
     c.real, c.imag = real, imag