matrix multiplication for complex numbers in PyTorch

Viewed 1697
2 Answers

Currently torch.matmul is not supported for complex tensors such as ComplexFloatTensor but you could do something as compact as the following code:

def matmul_complex(t1,t2):
    return torch.view_as_complex(torch.stack((t1.real @ t2.real - t1.imag @ t2.imag, t1.real @ t2.imag + t1.imag @ t2.real),dim=2))

When possible avoid using for loops as these will result in much slower implementations. Vectorization is achieved by using built-in methods as demonstrated in the code I have attached. For example, your code takes roughly 6.1s on CPU while the vectorized version takes only 101ms (~60 times faster) for 2 random complex matrices with dimensions 1000 X 1000.

Update:

Since PyTorch 1.7.0 (as @EduardoReis mentioned) you can do matrix multiplication between complex matrices similarly to real-valued matrices as follows:

t1 @ t2 (for t1, t2 complex matrices).

I implemented this function for pytorch.matmul for complex numbers using torch.mv and it's working fine for time-being:

def matmul_complex(t1, t2):
  m = list(t1.size())[0]
  n = list(t2.size())[1]
  t = torch.empty((1,n), dtype=torch.cfloat)
  t_total = torch.empty((m,n), dtype=torch.cfloat)
  for i in range(0,n):
    if i == 0:
      t_total = torch.mv(t1,t2[:,i])
    else:
      t_total = torch.cat((t_total, torch.mv(t1,t2[:,i])), 0)
  t_final = torch.reshape(t_total, (m,n))
  return t_final

I am new to PyTorch, so please correct me if I am wrong.

Related