1 #define TORCH_ASSERT_NO_OPERATORS 2 3 #include <ATen/native/UnaryOps.h> 4 5 #include <ATen/Dispatch.h> 6 #include <ATen/native/Math.h> 7 #include <ATen/native/TensorIterator.h> 8 #include <ATen/native/cpu/Loops.h> 9 10 namespace at::native { 11 inline namespace CPU_CAPABILITY { airy_ai_kernel(TensorIteratorBase & iterator)12static void airy_ai_kernel(TensorIteratorBase& iterator) { 13 TORCH_INTERNAL_ASSERT(iterator.ntensors() == 2); 14 15 AT_DISPATCH_FLOATING_TYPES(iterator.common_dtype(), "airy_ai_cpu", [&]() { 16 cpu_kernel(iterator, [](scalar_t x) { 17 return airy_ai_forward(x); 18 }); 19 }); 20 } // airy_ai_kernel(TensorIteratorBase& iterator) 21 } // namespace CPU_CAPABILITY 22 23 REGISTER_DISPATCH(special_airy_ai_stub, &CPU_CAPABILITY::airy_ai_kernel); 24 } // namespace at::native 25