1 #pragma once 2 3 #include <ATen/core/ATen_fwd.h> 4 #include <ATen/native/DispatchStub.h> 5 6 namespace at::native { 7 8 enum class NESTED_DENSE_OP : uint8_t { ADD, MUL }; 9 10 using nested_dense_elementwise_fn = void (*)( 11 Tensor& result, 12 const Tensor& self, 13 const Tensor& other, 14 const NESTED_DENSE_OP& op); 15 16 DECLARE_DISPATCH(nested_dense_elementwise_fn, nested_dense_elementwise_stub); 17 18 } // namespace at::native 19