Go to the documentation of this file.
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>
68 inline SwapAxisExp<SrcExp, DType, ExpInfo<SrcExp>::kDim,
73 a2 < a1>::Error_Expression_Does_Not_Meet_Dimension_Req();
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)),
85 shapez_(e.shape_[a1]),
86 shapec_(e.shape_.ProdShape(a2 + 1, a1)),
87 shapen_(e.shape_[a2]) {}
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)),
112 shapez_(e.shape_[a2]) {}
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
Plan(const SwapAxisExp< SrcExp, DType, dimsrc, m_a1, a2 > &e)
Definition: swapaxis.h:82
swap two axis of a tensor input: Tensor<Device,dim>: ishape output: Tensor<Device,...
Definition: swapaxis.h:44
used to help static type check
Definition: expr_engine-inl.h:330
MSHADOW_XINLINE DType Eval(index_t i, index_t x) const
Definition: swapaxis.h:113
#define MSHADOW_XINLINE
Definition: base.h:228
const SrcExp & src_
source expression
Definition: swapaxis.h:50
SwapAxisExp(const SrcExp &src)
constructor
Definition: swapaxis.h:52
static const int a1
Definition: swapaxis.h:48
static Shape< dim > Check(const E &t)
static type inference template, used to get the dimension of each expression, if ExpInfo<E>::kDim == ...
Definition: expr_engine-inl.h:262
MSHADOW_XINLINE DType Eval(index_t i, index_t j) const
Definition: swapaxis.h:88
Plan< BinaryMapExp< OP, TA, TB, DType, etype >, DType > MakePlan(const BinaryMapExp< OP, TA, TB, DType, etype > &e)
Definition: expr_engine-inl.h:239
Shape< dim > shape_
the shape of this expression
Definition: expr_engine-inl.h:47
static const int kDim
Definition: expr_engine-inl.h:263
int32_t index_t
type that will be used for index
Definition: base.h:328
Definition: expr_engine-inl.h:58
defines how expression exp can be evaluated and stored into dst
Definition: expression.h:79
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
overloaded + operator between half_t and bf16_t
Definition: base.h:319
a general class that allows extension that makes tensors of some shape
Definition: expr_engine-inl.h:43
Plan(const SwapAxisExp< SrcExp, DType, dimsrc, 1, a2 > &e)
Definition: swapaxis.h:108