ReduceOps.h 1.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657
  1. #pragma once
  2. #include <ATen/native/DispatchStub.h>
  3. #include <c10/util/ArrayRef.h>
  4. #include <optional>
  5. namespace c10 {
  6. class Scalar;
  7. }
  8. namespace at {
  9. struct TensorIterator;
  10. class Tensor;
  11. }
  12. namespace at::native {
  13. using reduce_fn = void(*)(TensorIterator &);
  14. DECLARE_DISPATCH(reduce_fn, sum_stub)
  15. DECLARE_DISPATCH(reduce_fn, nansum_stub)
  16. DECLARE_DISPATCH(reduce_fn, prod_stub)
  17. DECLARE_DISPATCH(reduce_fn, mean_stub)
  18. DECLARE_DISPATCH(reduce_fn, and_stub)
  19. DECLARE_DISPATCH(reduce_fn, or_stub)
  20. DECLARE_DISPATCH(reduce_fn, min_values_stub)
  21. DECLARE_DISPATCH(reduce_fn, max_values_stub)
  22. DECLARE_DISPATCH(reduce_fn, argmax_stub)
  23. DECLARE_DISPATCH(reduce_fn, argmin_stub)
  24. DECLARE_DISPATCH(reduce_fn, xor_sum_stub)
  25. using reduce_std_var_function =
  26. void (*)(TensorIterator&, double correction, bool take_sqrt);
  27. DECLARE_DISPATCH(reduce_std_var_function, std_var_stub)
  28. using reduce_norm_fn =
  29. void (*)(Tensor&, const Tensor&, const c10::Scalar&, std::optional<int64_t>);
  30. DECLARE_DISPATCH(reduce_norm_fn, norm_kernel)
  31. using reduce_fn_flag = void(*)(TensorIterator &, const c10::Scalar&);
  32. DECLARE_DISPATCH(reduce_fn_flag, norm_stub)
  33. using structured_cum_fn = void (*)(const Tensor&, const Tensor&, int64_t);
  34. using cum_fn = void (*)(Tensor&, const Tensor&, int64_t);
  35. DECLARE_DISPATCH(structured_cum_fn, cumsum_stub)
  36. DECLARE_DISPATCH(structured_cum_fn, cumprod_stub)
  37. DECLARE_DISPATCH(cum_fn, logcumsumexp_stub)
  38. DECLARE_DISPATCH(void (*)(const Tensor&, int64_t, bool, Tensor&, Tensor&), aminmax_stub)
  39. DECLARE_DISPATCH(void (*)(const Tensor&, Tensor&, Tensor&), aminmax_allreduce_stub)
  40. // Used in cuda/Normalization.cu
  41. TORCH_API std::tuple<Tensor&,Tensor&> var_mean_out(
  42. Tensor &result1, Tensor &result2, const Tensor &self, IntArrayRef dim,
  43. int64_t correction, bool keepdim);
  44. } // namespace at::native