File size: 963 Bytes
c1af2fa |
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 |
#pragma once
#include <ATen/native/DispatchStub.h>
#include <cstdint>
namespace at {
class Tensor;
namespace native {
using forward_fn = void (*)(const Tensor&, const Tensor&);
using backward_fn = void(*)(const Tensor &, const Tensor &, const Tensor&);
DECLARE_DISPATCH(forward_fn, softmax_lastdim_kernel)
DECLARE_DISPATCH(forward_fn, log_softmax_lastdim_kernel)
DECLARE_DISPATCH(backward_fn, softmax_backward_lastdim_kernel)
DECLARE_DISPATCH(backward_fn, log_softmax_backward_lastdim_kernel)
using forward_fn_with_dim = void(*)(const Tensor &, const Tensor &, const int64_t);
using backward_fn_with_dim =
void (*)(const Tensor&, const Tensor&, const Tensor&, const int64_t);
DECLARE_DISPATCH(forward_fn_with_dim, softmax_kernel)
DECLARE_DISPATCH(forward_fn_with_dim, log_softmax_kernel)
DECLARE_DISPATCH(backward_fn_with_dim, softmax_backward_kernel)
DECLARE_DISPATCH(backward_fn_with_dim, log_softmax_backward_kernel)
}
}
|