TensorModeKernel.h 431 B

12345678910111213141516171819
  1. #pragma once
  2. #include <cstdint>
  3. namespace at {
  4. class TensorBase;
  5. }
  6. namespace at {
  7. namespace native {
  8. void launch_fused_mode_kernel(
  9. const TensorBase &values, const TensorBase &indices,
  10. const TensorBase &self, int64_t slice_size, int64_t slices);
  11. void launch_apply_mode_kernel(
  12. const TensorBase &values, const TensorBase &indices,
  13. const TensorBase &self, int64_t dim, int64_t ndim);
  14. }} // namespace at::native