Pytorch has support for complex tensors and can also do Autograd for them.

https://pytorch.org/docs/stable/complex_numbers.html

https://pytorch.org/docs/stable/notes/autograd.html#autograd-for-complex-numbers

We can gate this behind a feature flag if we want. But it would be a nice to have.