[WIP]: Add polar method for complex tensors. #35563
Conversation
💊 Build failures summary and remediationsAs of commit 5adb0b5 (more details on the Dr. CI page):
🕵️ 2 new failures recognized by patternsThe following build failures do not appear to be due to upstream breakages:
|
| @@ -129,6 +129,18 @@ template <> class Vec256<std::complex<double>> { | |||
| auto angle = _mm256_permute_pd(angle_(), 0x05); // angle 90-angle | |||
| return _mm256_and_pd(angle, real_mask); // angle 0 | |||
| } | |||
| __m256d polar_() const { | |||
There was a problem hiding this comment.
The polar_() method is the key take-away here. It is essentially a private method for the Vec256 class.
Now we have a decision to make:
- We could modify the existing
operator*andoperator/to convert to polar coordinates, perform the op, and then convert to back to cartesian coordinates. I think you will find that we waste too much time converting between polar and cart for a single math kernel. Note: You lose considerable performance by calling multiple AVX functions. - Therefore, we might need to create a new dtype
ComplexPolarFloatdtype so that we can keep the data in polar coordinates for multiple consecutive math ops. This will require significantly more work, but you are getting there.
Method 2 implies that we can have a different input dtype than the output dtype. PyTorch traditionally isn't setup to do this, but it has been attempted:
- People have tried to cast complex to float for the real(), imag(), abs(), angle() kernels. This is hard because its also changes the memory size.
- le kernel successfully casts the output byte data type to boolean (both are 1 byte in memory).
- #35524 is implementing the std::complex<T> data type internally in PyTorch. You could copy this code and create a std::complex_polar<T> C++ data type. Then you could do the same thing as the le_kernel and convert from std::complex<T> to std::complex_polar<T> when calling polar()
There was a problem hiding this comment.
In method 2 would we also need to add Vec256<std::complex_polar<T>>? I will look into le kernel, thanks for link.
| template <typename TYPE, std::enable_if_t<!c10::is_complex_t<TYPE>::value, int> = 0> | ||
| inline std::complex<TYPE> polar_impl (TYPE a) { | ||
| return std::complex<TYPE> (std::abs(a), std::arg(a)); | ||
| } | ||
|
|
||
| template <typename TYPE, std::enable_if_t<c10::is_complex_t<TYPE>::value, int> = 0> | ||
| inline TYPE polar_impl (TYPE a) { | ||
| return TYPE(std::abs(a), std::arg(a)); | ||
| } | ||
|
|
There was a problem hiding this comment.
You might be able to use std::polar here, however there is no std::cart so you will need to do something like this anyways.
|
@kostekIV, @anjali411, @ezyang LGTM. We should discuss the possibility of a |
| Vec256<T> polar() const { | ||
| return *this; | ||
| } |
There was a problem hiding this comment.
Yes, in a template class you need to re-implement all methods during specialization`.
I think the default case should return abs(x), not real(x).
If you are publicly adding a new kernel, you also need to specify the derivate so that Autograd works in derivatives.yaml. You would have to combine the derivative of abs() and angle() in some way.
There was a problem hiding this comment.
Oh, yes you're right. I didn't know about derivates.yaml, will look into it.
| (0.0899 + 0.3232j), | ||
| (0.9718 + 0.1947j), | ||
| (0.2349 + 0.1463j)], dtype=torch.complex64) | ||
| >>> a.polar() |
There was a problem hiding this comment.
Does it make sense to represent polar representation of complex as a + bi?
There was a problem hiding this comment.
Probably not, wasn't sure how to store it differently in complex tensor, I guess it may be resolved by ComplexPolarFloat dtype.
|
@dylanbespalko I understand next steps to improve this solution are first create |
|
Correct, you will likely need to submit a proposal to create:
The use cases would be:
There will need to be a discussion because this could lead to Cylindrical coords, Spherical Coords, etc.. FYI, #11641 contains partial information on how the ComplexFloat dtype was added. |
|
By proposal do you mean pull request or issue? |
|
Sorry, I meant you should create a feature request. There are people at Fb that need to approve something with that large of a scope. https://pytorch.org/docs/stable/community/contribution_guide.html If I remember correctly, you create a new issue on GitHub and label as a Feature Request. Github should automatically apply this template The people that made PyTorch will discuss it with you. They are usually the people that make big changes. I would also learn how to benchmark kernels https://github.com/pytorch/pytorch/tree/master/benchmarks/operator_benchmark Perhaps you could modify the
|
|
@dylanbespalko lmk if you need help from me reviewing this |
This seems reasonable, though whether or not we should invest in writing it (writing one of these classes is no small effort) depends on how serious we're going to go about polar. |
|
Looks like you have some CI failures specifically in the doc tests. Try running test/test_doc_coverage.py on your local machine. Some of the CI machines also failed to connect, so add another commit to restart the CI. |
|
It looks like CI is failing on windows_vs2019 FAIL [0.004s]: test_polar_cpu_complex128 (__main__.TestTorchDeviceTypeCPU)
----------------------------------------------------------------------
Traceback (most recent call last):
File "C:\Users\circleci\project\build\win_tmp\build\torch\testing\_internal\common_device_type.py", line 224, in instantiated_test
result = test(self, device_arg, dtype)
File "C:\Users\circleci\project\build\win_tmp\build\torch\testing\_internal\common_device_type.py", line 431, in only_fn
return fn(slf, device, *args, **kwargs)
File "test_torch.py", line 15175, in test_polar
self.assertEqual(polar_tensor, a.abs() + 1j * a.angle())
File "C:\Users\circleci\project\build\win_tmp\build\torch\testing\_internal\common_utils.py", line 894, in assertEqual
assertTensorsEqual(x, y)
File "C:\Users\circleci\project\build\win_tmp\build\torch\testing\_internal\common_utils.py", line 860, in assertTensorsEqual
self.assertLessEqual(max_err, prec, message)
AssertionError: tensor(1.5707, dtype=torch.float64) not less than or equal to 1e-05 :
======================================================================
FAIL [0.004s]: test_polar_cpu_complex64 (__main__.TestTorchDeviceTypeCPU)
----------------------------------------------------------------------
Traceback (most recent call last):
File "C:\Users\circleci\project\build\win_tmp\build\torch\testing\_internal\common_device_type.py", line 224, in instantiated_test
result = test(self, device_arg, dtype)
File "C:\Users\circleci\project\build\win_tmp\build\torch\testing\_internal\common_device_type.py", line 431, in only_fn
return fn(slf, device, *args, **kwargs)
File "test_torch.py", line 15175, in test_polar
self.assertEqual(polar_tensor, a.abs() + 1j * a.angle())
File "C:\Users\circleci\project\build\win_tmp\build\torch\testing\_internal\common_utils.py", line 894, in assertEqual
assertTensorsEqual(x, y)
File "C:\Users\circleci\project\build\win_tmp\build\torch\testing\_internal\common_utils.py", line 860, in assertTensorsEqual
self.assertLessEqual(max_err, prec, message)
AssertionError: tensor(1.5708) not less than or equal to 1e-05 :
----------------------------------------------------------------------
Ran 4781 tests in 358.397s
FAILED (failures=2, skipped=201)
|
I am looking into it, it might take some time for me as I dont have machine with windows, and on mac or linux I didn't succeed in reproducing it yet. |
|
The test case looks really good. I'm a little nervous about I also found out that they want to isolate the complex number test cases in I have had Windows failures int the past. I fix them by making changes and re-running CI. Not elegant, but it happens a lot. |
|
I've added #36029 to track any issues with complex unit testing. |
|
@dylanbespalko @anjali411 Sorry for lack of word/work past few weeks. I am wondering what I should do with this pr now. Since polar complex most likely won't be a thing in near future should I proceed with |
Test FailuresIn you test case, you have a line Other things to try:
Is
|
|
My apologies, I don't think you can simply re-use the cartesian functions. They need to be redefined for polar coordinates. It's really hard to get these things going. |
|
Hi @kostekIV! Thank you for your pull request and welcome to our community. Action RequiredIn order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you. ProcessIn order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA. Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with If you have received this in error or have any questions, please contact us at [email protected]. Thanks! |
|
Looks like this PR hasn't been updated in a while so we're going to go ahead and mark this as |
Hi @dylanbespalko its still work in progress but I would like to ask for your opinion. I was trying to follow your instruction and started with method
polar. So from input it is producing its polar form. Currently magnitude is stored as real value and angle as imaginary. First I would like to ask you if this is like you imagined it should be? Some problems that I run into and not sure how to tackle them. Polar form makes sense not only for complex type tensor but also for other like float or double but it is hard implement it because I understand it like that if tensor is computed by kernel function its dtype will match dtype of tensor that was used as input to kernel. I was thinking of casting such tensors to complex in fileaten/src/ATen/native/UnaryOps.cppbefore passing them to polar implementations. Also I had to add function polar invec256_base.hwhich is just a placeholder to make it compile, not sure how to do it better here.Related to #35312