/usr/local/lib64/python3.6/site-packages/torch/include/ATen/native
Edit: /usr/local/lib64/python3.6/site-packages/torch/include/ATen/native/SegmentReduce.h (685B)
#pragma once
#include
#include
#include
namespace at {
namespace native {
enum SegmentReductionType { MAX, MEAN, MIN, SUM };
using segment_reduce_fn = Tensor (*)(
SegmentReductionType,
const Tensor&,
const Tensor&,
int64_t,
const c10::optional&);
DECLARE_DISPATCH(segment_reduce_fn, _segment_reduce_stub);
using segment_reduce_backward_fn = Tensor (*)(
const Tensor&,
const Tensor&,
const Tensor&,
SegmentReductionType,
const Tensor&,
int64_t);
DECLARE_DISPATCH(segment_reduce_backward_fn, _segment_reduce_backward_stub);
} // namespace native
} // namespace at