1 #define TORCH_ASSERT_ONLY_METHOD_OPERATORS
2 #include <ATen/Dispatch.h>
3 #include <ATen/native/TensorFactories.h>
4 #include <ATen/native/TensorIterator.h>
5 #include <ATen/native/cpu/Loops.h>
6
7 namespace at::native {
8 namespace {
9
complex_kernel(TensorIterator & iter)10 void complex_kernel(TensorIterator& iter) {
11 AT_DISPATCH_FLOATING_TYPES_AND(kHalf, iter.input_dtype(), "complex_cpu", [&]() {
12 cpu_kernel(iter, [=](scalar_t a, scalar_t b) -> c10::complex<scalar_t> {
13 return c10::complex<scalar_t>(a, b);
14 });
15 });
16 }
17
polar_kernel(TensorIterator & iter)18 void polar_kernel(TensorIterator& iter) {
19 AT_DISPATCH_FLOATING_TYPES(iter.input_dtype(), "polar_cpu", [&]() {
20 cpu_kernel(iter, [=](scalar_t a, scalar_t b) -> c10::complex<scalar_t> {
21 return c10::complex<scalar_t>(a * std::cos(b), a * std::sin(b));
22 });
23 });
24 }
25
26 } // anonymous namespace
27
28 REGISTER_DISPATCH(complex_stub, &complex_kernel);
29 ALSO_REGISTER_AVX512_DISPATCH(polar_stub, &polar_kernel);
30
31 } // namespace at::native
32