26 #ifndef MSHADOW_EXTENSION_SWAPAXIS_H_ 27 #define MSHADOW_EXTENSION_SWAPAXIS_H_ 30 #include "../extension.h" 44 template<
typename SrcExp,
typename DType,
int dimsrc,
int m_a1,
int a2>
46 public MakeTensorExp<SwapAxisExp<SrcExp, DType, dimsrc, m_a1, a2>,
47 SrcExp, dimsrc, DType> {
49 static const int a1 = dimsrc - m_a1;
68 template<
int a1,
int a2,
typename SrcExp,
typename DType,
int etype>
74 a2 < a1>::Error_Expression_Does_Not_Meet_Dimension_Req();
75 return SwapAxisExp<SrcExp, DType, ExpInfo<SrcExp>::kDim,
78 template<
typename SrcExp,
typename DType,
int dimsrc,
int m_a1,
int a2>
82 static const int a1 = dimsrc - m_a1;
85 shapey_(e.
shape_.ProdShape(a1 + 1, dimsrc - 1)),
87 shapec_(e.
shape_.ProdShape(a2 + 1, a1)),
98 return src_.Eval(((((i / shapen_) * shapez_ + z) * shapec_ +
99 c) * shapen_ + n) * shapey_ + y, j);
104 const index_t shapey_, shapez_, shapec_, shapen_;
106 template<
typename SrcExp,
typename DType,
int dimsrc,
int a2>
111 shapex_(e.
shape_[dimsrc - 1]),
112 shapey_(e.
shape_.ProdShape(a2 + 1, dimsrc - 1)),
120 return src_.Eval((n * shapex_ + x) * shapey_ + y , z);
125 const index_t shapex_, shapey_, shapez_;
129 #endif // MSHADOW_EXTENSION_SWAPAXIS_H_ Definition: expr_engine-inl.h:59
used to help static type check
Definition: expr_engine-inl.h:331
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:71
static Shape< dim > Check(const E &t)
MSHADOW_XINLINE DType Eval(index_t i, index_t j) const
Definition: swapaxis.h:89
#define MSHADOW_XINLINE
Definition: base.h:223
static type inference template, used to get the dimension of each expression, if ExpInfo<E>::kDim == ...
Definition: expr_engine-inl.h:263
Plan(const SwapAxisExp< SrcExp, DType, dimsrc, m_a1, a2 > &e)
Definition: swapaxis.h:83
int32_t index_t
type that will be used for index
Definition: base.h:336
const SrcExp & src_
source expression
Definition: swapaxis.h:51
static const int a1
Definition: swapaxis.h:49
Plan(const SwapAxisExp< SrcExp, DType, dimsrc, 1, a2 > &e)
Definition: swapaxis.h:109
defines how expression exp can be evaluated and stored into dst
Definition: expression.h:80
const SubType & self(void) const
Definition: expression.h:83
Plan< BinaryMapExp< OP, TA, TB, DType, etype >, DType > MakePlan(const BinaryMapExp< OP, TA, TB, DType, etype > &e)
Definition: expr_engine-inl.h:240
a general class that allows extension that makes tensors of some shape
Definition: expr_engine-inl.h:44
overloaded + operator between half_t and bf16_t
Definition: base.h:327
Shape< dim > shape_
the shape of this expression
Definition: expr_engine-inl.h:48
SwapAxisExp(const SrcExp &src)
constructor
Definition: swapaxis.h:53
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:45
MSHADOW_XINLINE DType Eval(index_t i, index_t x) const
Definition: swapaxis.h:114