1 #pragma once 2 #include <ATen/native/DispatchStub.h> 3 4 namespace at { 5 class TensorBase; 6 } 7 8 namespace at::native { 9 10 using pixel_shuffle_fn = void(*)(TensorBase&, const TensorBase&, int64_t); 11 DECLARE_DISPATCH(pixel_shuffle_fn, pixel_shuffle_kernel); 12 DECLARE_DISPATCH(pixel_shuffle_fn, pixel_unshuffle_kernel); 13 14 } // at::native 15