25 #ifndef MSHADOW_EXTENSION_REDUCE_WITH_AXIS_H_ 26 #define MSHADOW_EXTENSION_REDUCE_WITH_AXIS_H_ 28 #include "../extension.h" 38 template<
typename Reducer,
typename SrcExp,
typename DType,
int dimsrc,
bool mask,
int dimdst>
40 public MakeTensorExp<ReduceWithAxisExp<Reducer, SrcExp, DType, dimsrc, mask, dimdst>,
41 SrcExp, dimdst, DType> {
55 bool keepdim = (dimsrc == dimdst);
56 CHECK(dimsrc > axis) <<
"reduce axis out of bound";
58 for (
int i = 0; i < axis; ++i) {
59 this->
shape_[i] = src_shape[i];
61 this->size_ = src_shape[axis];
64 for (
int i = axis + 1; i < dimsrc; ++i) {
65 this->trailing_ *= src_shape[i];
66 this->
shape_[i - 1] = src_shape[i];
70 for (
index_t i = axis + 1; i < dimsrc; ++i) {
71 this->trailing_ *= src_shape[i];
72 this->
shape_[i] = src_shape[i];
76 this->last_ = src_shape[dimsrc - 1];
77 this->last_dst_dim_ = this->
shape_[dimdst - 1];
89 template<
typename Reducer,
bool mask,
typename SrcExp,
typename DType,
int etype>
93 return ReduceWithAxisExp<Reducer, SrcExp, DType, ExpInfo<SrcExp>::kDim,
mask,
105 template<
typename Reducer,
bool mask,
typename SrcExp,
typename DType,
int etype>
106 inline ReduceWithAxisExp<Reducer, SrcExp, DType, ExpInfo<SrcExp>::kDim,
mask,
109 return ReduceWithAxisExp<Reducer, SrcExp, DType, ExpInfo<SrcExp>::kDim,
mask,
116 template<
typename Reducer,
typename SrcExp,
typename DType,
int dimsrc,
bool mask,
int dimdst>
128 DType res; Reducer::SetInitValue(res);
137 return static_cast<DType
>(
static_cast<int>(idx));
139 DType res; Reducer::SetInitValue(res);
154 #endif // MSHADOW_EXTENSION_REDUCE_WITH_AXIS_H_ index_t last_
size of last src dimension
Definition: reduce_with_axis.h:51
const SubType & self(void) const
Definition: expression.h:82
Definition: expr_engine-inl.h:58
static Shape< dim > Check(const E &t)
#define MSHADOW_XINLINE
Definition: base.h:230
static type inference template, used to get the dimension of each expression, if ExpInfo<E>::kDim == ...
Definition: expr_engine-inl.h:262
MSHADOW_XINLINE bool IsNan(volatile DType val)
Definition: base.h:682
MaskExp< IndexExp, SrcExp, DType > mask(const Exp< IndexExp, DType, e1 > &index, const Exp< SrcExp, DType, e2 > &src)
Definition: mask.h:57
const SrcExp & src_
source oprand
Definition: reduce_with_axis.h:43
int32_t index_t
type that will be used for index
Definition: base.h:343
index_t size_
size of axis dimension
Definition: reduce_with_axis.h:49
MSHADOW_XINLINE DType Eval(index_t i, index_t j) const
Definition: reduce_with_axis.h:122
ReduceWithAxisExp< Reducer, SrcExp, DType, ExpInfo< SrcExp >::kDim, mask, ExpInfo< SrcExp >::kDim - 1 > reduce_with_axis(const Exp< SrcExp, DType, etype > &src, int axis)
reduce out the dimension of src labeled by axis.
Definition: reduce_with_axis.h:92
ReduceWithAxisExp(const SrcExp &src, int axis)
Definition: reduce_with_axis.h:53
reduce out the dimension of src labeled by axis.
Definition: reduce_with_axis.h:39
defines how expression exp can be evaluated and stored into dst
Definition: expression.h:79
Plan< BinaryMapExp< OP, TA, TB, DType, etype >, DType > MakePlan(const BinaryMapExp< OP, TA, TB, DType, etype > &e)
Definition: expr_engine-inl.h:239
ReduceWithAxisExp< Reducer, SrcExp, DType, ExpInfo< SrcExp >::kDim, mask, ExpInfo< SrcExp >::kDim > reduce_keepdim(const Exp< SrcExp, DType, etype > &src, int axis)
reduce out the dimension of src labeled by axis, keepdim turned on.
Definition: reduce_with_axis.h:108
a general class that allows extension that makes tensors of some shape
Definition: expr_engine-inl.h:43
overloaded + operator between half_t and bf16_t
Definition: base.h:334
Shape< dim > shape_
the shape of this expression
Definition: expr_engine-inl.h:47
Plan(const ReduceWithAxisExp< Reducer, SrcExp, DType, dimsrc, mask, dimdst > &e)
Definition: reduce_with_axis.h:119
index_t trailing_
size of trailing dimensions
Definition: reduce_with_axis.h:47
index_t last_dst_dim_
size of last destination dimension
Definition: reduce_with_axis.h:45