Why does the square of a complex tensor with no imaginary part yields a result with an imaginary part?

I have a complex tensor: tensor1 = tensor([0.0000+0.j, -106990.0794+0.j], device='cuda:1', dtype=torch.complex128)

The both elements have no imaginary part (0.j), however, when squaring the variable tensor1**2 it yields an imaginary part in the second element

tensor([0.0000e+00+0.0000e+00j, 1.1447e+10-2.8037e-06j], device='cuda:1', dtype=torch.complex128)

This behavior is not exhibited when performing the square in numpy

tensor1.cpu().numpy()**2

Out : array([0.00000000e+00+0.j, 1.14468771e+10-0.j])

Why are they different?

Solution

This looks to be a numeric issue impacting complex numbers on GPU. Weirdly, it impacts negative numbers but not positive.

values = [5.0, -5.0]
devices = ['cpu', 'cuda']

for device in devices:
    for value in values:
        x = torch.tensor(value+0.j, device=device, dtype=torch.complex128)
        x_pow2 = x.pow(2)
        x_x_conj = x * x.conj()
        check1 = x_pow2 == x_x_conj
        check2 = x_pow2.imag == 0
        check3 = x_x_conj.imag==0
        
        print(f"{device}\t{value}\t{check1.item()}\t{check2.item()}\t{check3.item()}")

cpu 5.0 True    True    True
cpu -5.0    True    True    True
cuda    5.0 False   True    True
cuda    -5.0    False   False   True

For a given input x = value + 0.j, the code computes x_pow2 = x.pow(2) and x_x_conj = x * x.conj(). The code checks if x_pow2 == x_x_conj, and if the imaginary component of those tensors is zero.

We see on CPU, x_pow2 == x_x_conj for both input values, and both x_pow2 and x_x_conj have a zero imaginary component.

On GPU, we see that in both cases x_pow2 != x_x_conj. For the case of value = 5.0, both x_pow2 and x_x_conj have zero imaginary component. For the case of value = -5.0, x_x_conj has zero imaginary component but x_pow2 does not.

I've tested this with a few values, and I consistently see inputs with positive real values yielding x_pow2 with zero imaginary component, while inputs with negative real values yielding x_pow2 with nonzero imaginary component.

This is likely due to something weird happening in the pow kernel that isn't present in the mul kernel.