pytorch/aten/src/ATen/native/cpu/GridSamplerKernel.h at main · gnosyslambda/pytorch · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
#pragma once

#include <ATen/native/DispatchStub.h>

#include <array>
#include <cstdint>

namespace at {
class TensorBase;
}

namespace at::native {

using forward_2d_fn = void (*) (
    const TensorBase &output,
    const TensorBase &input,
    const TensorBase &grid,
    int64_t interpolation_mode,
    int64_t padding_mode,
    bool align_corners);
using backward_2d_fn = void (*) (
    const TensorBase &grad_input,
    const TensorBase &grad_grid,
    const TensorBase &grad_output,
    const TensorBase &input,
    const TensorBase &grid,
    int64_t interpolation_mode,
    int64_t padding_mode,
    bool align_corners,
    std::array<bool, 2> output_mask);
DECLARE_DISPATCH(forward_2d_fn, grid_sampler_2d_cpu_kernel)
DECLARE_DISPATCH(backward_2d_fn, grid_sampler_2d_backward_cpu_kernel)

} // namespace at::native