25 #ifndef MSHADOW_EXTENSION_SWAPAXIS_H_ 26 #define MSHADOW_EXTENSION_SWAPAXIS_H_ 29 #include "../extension.h" 43 template<
typename SrcExp,
typename DType,
int dimsrc,
int m_a1,
int a2>
45 public MakeTensorExp<SwapAxisExp<SrcExp, DType, dimsrc, m_a1, a2>,
46 SrcExp, dimsrc, DType> {
48 static const int a1 = dimsrc - m_a1;
67 template<
int a1,
int a2,
typename SrcExp,
typename DType,
int etype>
73 a2 < a1>::Error_Expression_Does_Not_Meet_Dimension_Req();
74 return SwapAxisExp<SrcExp, DType, ExpInfo<SrcExp>::kDim,
77 template<
typename SrcExp,
typename DType,
int dimsrc,
int m_a1,
int a2>
81 static const int a1 = dimsrc - m_a1;
84 shapey_(e.
shape_.ProdShape(a1 + 1, dimsrc - 1)),
86 shapec_(e.
shape_.ProdShape(a2 + 1, a1)),
97 return src_.Eval(((((i / shapen_) * shapez_ + z) * shapec_ +
98 c) * shapen_ + n) * shapey_ + y, j);
103 const index_t shapey_, shapez_, shapec_, shapen_;
105 template<
typename SrcExp,
typename DType,
int dimsrc,
int a2>
110 shapex_(e.
shape_[dimsrc - 1]),
111 shapey_(e.
shape_.ProdShape(a2 + 1, dimsrc - 1)),
119 return src_.Eval((n * shapex_ + x) * shapey_ + y , z);
124 const index_t shapex_, shapey_, shapez_;
128 #endif // MSHADOW_EXTENSION_SWAPAXIS_H_ const SubType & self(void) const
Definition: expression.h:82
Definition: expr_engine-inl.h:58
used to help static type check
Definition: expr_engine-inl.h:330
SwapAxisExp< SrcExp, DType, ExpInfo< SrcExp >::kDim, ExpInfo< SrcExp >::kDim - a1, a2 > swapaxis(const Exp< SrcExp, DType, etype > &src)
a expression that reshapes a tensor to another shape
Definition: swapaxis.h:70
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
Plan(const SwapAxisExp< SrcExp, DType, dimsrc, m_a1, a2 > &e)
Definition: swapaxis.h:82
int32_t index_t
type that will be used for index
Definition: base.h:343
const SrcExp & src_
source expression
Definition: swapaxis.h:50
static const int a1
Definition: swapaxis.h:48
MSHADOW_XINLINE DType Eval(index_t i, index_t j) const
Definition: swapaxis.h:88
MSHADOW_XINLINE DType Eval(index_t i, index_t x) const
Definition: swapaxis.h:113
Plan(const SwapAxisExp< SrcExp, DType, dimsrc, 1, a2 > &e)
Definition: swapaxis.h:108
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
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
SwapAxisExp(const SrcExp &src)
constructor
Definition: swapaxis.h:52
swap two axis of a tensor input: Tensor<Device,dim>: ishape output: Tensor<Device,dimdst> oshape[a1],oshape[a2] = ishape[a2],oshape[a1]
Definition: swapaxis.h:44