import torch


def complex_multiplication(a, b):
  """ Computes element-wise complex-tensor multiplication.

  Args:
    a and b: multiplication operands (tensors of complex numbers). Each
      tensor is in (batch x channel x height x width x complex_channel) format,
      where complex_channel contains the real and imaginary parts of each
      complex number in the tensor.

  Returns:
    results of a x b = (c+di)(j+hi) = (cj - dh) + (jd + ch)i.
  """

  assert (len(a.shape) >= 4 and len(a.shape) >= 4)
  assert (a.shape[-3:] == b.shape[-3:])
  assert (a.shape[-1] == 2)
  if len(a.shape) == 4:
    a = torch.unsqueeze(a, dim=1)
  if len(b.shape) == 4:
    b = torch.unsqueeze(b, dim=1)
  real_a = a[:, :, :, :, 0]
  imag_a = a[:, :, :, :, 1]
  real_b = b[:, :, :, :, 0]
  imag_b = b[:, :, :, :, 1]
  result = torch.stack([real_a * real_b - imag_a * imag_b,
                        real_a * imag_b + imag_a * real_b], dim=-1)
  return result