Hello, I would like to work on this, can I? If so, what is a good place for this functionality?

@kostekIV,

Here is a short range vs. long range solution to get you into PyTorch development. I've added some links to some relevant files you will need to change.

Short Range Solution: Add a polar() and cart() math kernels.

1. Learn how to register a new math kernel in PyTorch.
   • aten/src/ATen/native/native_functions.yaml
2. Implement the non-vectorized math kernels

• aten/src/ATen/native/cpu/UnaryOpsKernel.cpp
• See how sinh_kernel is implemented.

3. Implement the vectorized math kernels

• aten/src/ATen/native/cpu/UnaryOpsKernel.cpp
• See how neg_kernel is implemented.
• aten/src/ATen/cpu/vec256/vec256_complex_float.h
• aten/src/ATen/cpu/vec256/vec256_complex_double.h

Long Range Solution: Add ComplexPolarFloat dtype

The short range solution does not provide PyTorch with the knowledge that your complex number is in polar or cartesian coordinates. Hence all operators currently assume cartesian coordinates. We have started to discuss adding a Complex Polar dtype.

1. Work with fb to add a new polar dtypes (ComplexPolarHalf, ComplexPolarFloat, ComplexPolarDouble).
2. Add math kernels to the ComplexPolar Dtype to solve this issue.
   • Vectorized: Create a new vec256 file aten/src/ATen/cpu/vec256/vec256_complex_polar float.h

I would just focus on the short-range solution for now, but I think we need the long-range solution to speed up some of these computations.

@dylanbespalko
Thank you for such a comprehensive explanation. I will start with short range solution.

For reference, Numpy doesn't seem to have this yet: numpy/numpy#5228

@ezyang
Correct, numpy does not support polar() or cart() and perhaps these kernels should stay undocumented in the short-range solution.

The long-range solution of creating a ComplexPolarFloat dtype will likely run into a problem where the cast operation ComplexFloat.to(ComplexPolarFloat) performs a Scalar cast that does not benefit from vectorization.

Hence, the end-goal here should be to implement a vectorized polar() and cart() private method that is only used inside the Vec256 class when ComplexFloat or ComplexPolarFloat has to change coordinate systems.

@kostekIV

#35524 contains some new files that you should be aware of.