@anjali411 @ezyang Hi, I saw the update about having the polar keyword, but was wondering, is something like what I wrote along the right track to implement torch.complex?

💊 CI failures summary and remediations

As of commit a481888 (more details on the Dr. CI page):

💚 💚 Looks good so far! There are no failures yet. 💚 💚

Hi @anjali411 @mruberry, made some updates. A couple question/comments,

• Right now Vec256 isn't supported, but I'm not sure if it's possible since a real and imaginary Vec256<float> will lead to two Vec256<c10::complex<float>>'s.
• How could I go about supporting the function for CUDA?
• For the derivatives, I saw a lot of entries in derivatives.yaml don't use a _backward kernel. I'm thinking about just adding an entry like

- name: complex(Tensor real, Tensor imag) -> Tensor
  real: grad
  imag: grad * c10::complex<float>{0.0, 1.0}

- name: complex_polar(Tensor abs, Tensor angle) -> Tensor
  abs: grad * (angle.cos() + c10::complex<float>{0.0, 1.0} * angle.sin())
  angle: grad * abs * (- angle.sin() + c10::complex<float>{0.0, 1.0} * angle.cos())

However, would it be better to have specific complex_backward and complex_polar_backward methods that do a more optimized version of this? (essentially same way how complex and complex_polar kernels are implemented right now)

Hi @anjali411 @mruberry, made some updates. A couple question/comments,

• Right now Vec256 isn't supported, but I'm not sure if it's possible since a real and imaginary Vec256<float> will lead to two Vec256<c10::complex<float>>'s.
• How could I go about supporting the function for CUDA?
• For the derivatives, I saw a lot of entries in derivatives.yaml don't use a _backward kernel. I'm thinking about just adding an entry like

- name: complex(Tensor real, Tensor imag) -> Tensor
  real: grad
  imag: grad * c10::complex<float>{0.0, 1.0}

- name: complex_polar(Tensor abs, Tensor angle) -> Tensor
  abs: grad * (angle.cos() + c10::complex<float>{0.0, 1.0} * angle.sin())
  angle: grad * abs * (- angle.sin() + c10::complex<float>{0.0, 1.0} * angle.cos())

However, would it be better to have specific complex_backward and complex_polar_backward methods that do a more optimized version of this? (essentially same way how complex and complex_polar kernels are implemented right now)

• It's OK to not support Vec256.
• You're halfway there. Check out the binary CUDA kernels that use TensorIterator in aten/src/ATen/native/cuda. For example

  pytorch/aten/src/ATen/native/cuda/BinaryAddSubKernel.cu

  Line 11 in fcc9a1e

  void add_kernel_cuda(TensorIterator& iter, Scalar alpha_scalar) {

  .
• I think it's OK to not write a custom backward kernel for now. Let's focus on correctness.

Also see: https://github.com/pytorch/pytorch/pull/40246/files, which might give you some ideas for manipulating the complex types.

@mruberry @anjali411 sorry for just getting back to this. I think backward won't work altogether since it's fitting a complex gradient into a float input tensor. This was making me wonder how this function should work. Should it be able to accept complex inputs as well (and just use their real components)? It also made me wonder if this function would really be useful, since my guess is once x and y are both complex, the speed advantage over x + 1j * y wouldn't be that great.

@mruberry @anjali411 sorry for just getting back to this. I think backward won't work altogether since it's fitting a complex gradient into a float input tensor. This was making me wonder how this function should work. Should it be able to accept complex inputs as well (and just use their real components)? It also made me wonder if this function would really be useful, since my guess is once x and y are both complex, the speed advantage over x + 1j * y wouldn't be that great.

@wjgan7 I think backward definition should be similar to that of view_as_complex ( cc. @albanD ).

The function should only be able to expect floating point tensors (cc. @mruberry). For cpu, you can define a function in this file: aten/src/ATen/native/cpu/zmath.h which can be passed to the cpu_kernel. The function would accept two floating points a and b and return a c10::complex<T>(a, b). check out this PR as an example: #39955

As Mike mentioned, let's skip Vec256 implementation for this PR.

For the backward formula, I think we can just wait on the formula Anjali is going to put in the doc and apply it here. That will contain all the real to complex/complex to real changes that need to be applied.

@anjali411 I believe I already have the non-Vec256 implementation in native/cpu/ComplexHelperKernel.cpp. I can move that to native/cpu/zmath.h though. Where do you think I should have the higher level functions? Currently they are in native/ComplexHelper.cpp. Oops, I misunderstood what you meant, will do.

@albanD Sounds good, I can just follow what @mruberry said about the CUDA implementation and write tests for that then.

@wjgan7 any updates?

@anjali411 Apologies for delays, will try to get an update soon.

@anjali411 Not sure if I can view the failing CircleCI tests. Do you know what the issue is?

@wjgan7 they are just upstream failures, but if you rebase on the latest master, CI should be green.

@anjali411 merged this pull request in e437308.

@anjali411 @mruberry Thanks for all the reviews, I learned a lot!