Atan2 pytorch onnx

atan2.py

def my_atan2(y, x):
    pi = torch.from_numpy(np.array([np.pi])).to(y.device, y.dtype)
    ans = torch.atan(y / (x + 1e-6))
    ans += ((y > 0) & (x < 0)) * pi
    ans -= ((y < 0) & (x < 0)) * pi
    ans *= (1 - ((y > 0) & (x == 0)) * 1.0)
    ans += ((y > 0) & (x == 0)) * (pi / 2)
    ans *= (1 - ((y < 0) & (x == 0)) * 1.0)
    ans += ((y < 0) & (x == 0)) * (-pi / 2)
    return ans

Atan2 PyTorch implementation that can be exported to onnx

Thanks a lot!!!!!!atan2 is not supported in onnx and i've been trapped for a long time....

Thanks! It works.

How does this write into the project?

Better add an epsilon when dividing by x:

ans = torch.atan(y / (x + 1e-6))

@kometa-triatlon no need, torch atan handles inf, eg:

a = torch.from_numpy(np.array([1.0]))
b = torch.from_numpy(np.array([0.0]))
torch.atan(a/b)

Result

tensor([1.5708], dtype=torch.float64)

There can be a combination of parameters that would surprise you ;)

zero = torch.tensor([0.])
one = torch.tensor([1.])

print(my_atan2(zero, one), torch.atan2(zero, one))
print(my_atan2(one, zero), torch.atan2(one, zero))
print(my_atan2(one, one), torch.atan2(one, one))
print(my_atan2(zero, zero), torch.atan2(zero, zero))

Result:

tensor([0.]) tensor([0.])
tensor([1.5708]) tensor([1.5708])
tensor([0.7854]) tensor([0.7854])
tensor([nan]) tensor([0.])

With adding epsilon (ans = torch.atan(y/ (x + 1e-10))):

tensor([0.]) tensor([0.])
tensor([1.5708]) tensor([1.5708])
tensor([0.7854]) tensor([0.7854])
tensor([0.]) tensor([0.])

But atan2 of 0,0 is undefined, not 0. I think it's better to return nan, that way you would know that something is wrong.
Edit: okay I see most libraries handle 0,0 as 0, I'll add the epsilon, thanks for the suggestion

Hi @nikola-j , thank you for sharing. Have you tried this implementation with complex Tensors? If possible, could you share how you derive the aboved algorithm?
Thank you in advance!

Edited: I found it here: https://en.wikipedia.org/wiki/Atan2. Thank you!!!