File size: 431 Bytes
9dd3461
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
#pragma once
#include <cstdint>

namespace at {
class TensorBase;
}

namespace at {
namespace native {

void launch_fused_mode_kernel(
    const TensorBase &values, const TensorBase &indices,
    const TensorBase &self, int64_t slice_size, int64_t slices);

void launch_apply_mode_kernel(
    const TensorBase &values, const TensorBase &indices,
    const TensorBase &self, int64_t dim, int64_t ndim);

}}  // namespace at::native