xref: /aosp_15_r20/external/pytorch/aten/src/ATen/native/cuda/TensorModeKernel.h (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1 #pragma once
2 #include <cstdint>
3 
4 namespace at {
5 class TensorBase;
6 }
7 
8 namespace at {
9 namespace native {
10 
11 void launch_fused_mode_kernel(
12     const TensorBase &values, const TensorBase &indices,
13     const TensorBase &self, int64_t slice_size, int64_t slices);
14 
15 void launch_apply_mode_kernel(
16     const TensorBase &values, const TensorBase &indices,
17     const TensorBase &self, int64_t dim, int64_t ndim);
18 
19 }}  // namespace at::native
20