mxnet
swapaxis.h
Go to the documentation of this file.
1 /*
2  * Licensed to the Apache Software Foundation (ASF) under one
3  * or more contributor license agreements. See the NOTICE file
4  * distributed with this work for additional information
5  * regarding copyright ownership. The ASF licenses this file
6  * to you under the Apache License, Version 2.0 (the
7  * "License"); you may not use this file except in compliance
8  * with the License. You may obtain a copy of the License at
9  *
10  * http://www.apache.org/licenses/LICENSE-2.0
11  *
12  * Unless required by applicable law or agreed to in writing,
13  * software distributed under the License is distributed on an
14  * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15  * KIND, either express or implied. See the License for the
16  * specific language governing permissions and limitations
17  * under the License.
18  */
19 
26 #ifndef MSHADOW_EXTENSION_SWAPAXIS_H_
27 #define MSHADOW_EXTENSION_SWAPAXIS_H_
28 #include <algorithm>
29 #include <utility>
30 #include "../extension.h"
31 namespace mshadow {
32 namespace expr {
44 template<typename SrcExp, typename DType, int dimsrc, int m_a1, int a2>
45 struct SwapAxisExp:
46  public MakeTensorExp<SwapAxisExp<SrcExp, DType, dimsrc, m_a1, a2>,
47  SrcExp, dimsrc, DType> {
48  // decode the a1, a2
49  static const int a1 = dimsrc - m_a1;
51  const SrcExp &src_;
53  explicit SwapAxisExp(const SrcExp &src) : src_(src) {
55  std::swap(this->shape_[a1], this->shape_[a2]);
56  }
57 };
68 template<int a1, int a2, typename SrcExp, typename DType, int etype>
72  typedef ExpInfo<SrcExp> Info;
73  TypeCheckPass<Info::kDim >= a1 + 1 && Info::kDim >= a2 + 1 &&
74  a2 < a1>::Error_Expression_Does_Not_Meet_Dimension_Req();
75  return SwapAxisExp<SrcExp, DType, ExpInfo<SrcExp>::kDim,
76  ExpInfo<SrcExp>::kDim - a1, a2>(src.self());
77 }
78 template<typename SrcExp, typename DType, int dimsrc, int m_a1, int a2>
79 struct Plan<SwapAxisExp<SrcExp, DType, dimsrc, m_a1, a2>, DType> {
80  public:
81  // decode the a1
82  static const int a1 = dimsrc - m_a1;
84  : src_(MakePlan(e.src_)),
85  shapey_(e.shape_.ProdShape(a1 + 1, dimsrc - 1)),
86  shapez_(e.shape_[a1]),
87  shapec_(e.shape_.ProdShape(a2 + 1, a1)),
88  shapen_(e.shape_[a2]) {}
89  MSHADOW_XINLINE DType Eval(index_t i, index_t j) const {
90  const index_t y = i % shapey_;
91  i /= shapey_;
92  const index_t z = i % shapez_;
93  i /= shapez_;
94  const index_t c = i % shapec_;
95  i /= shapec_;
96  const index_t n = i % shapen_;
97  // swap z and n
98  return src_.Eval(((((i / shapen_) * shapez_ + z) * shapec_ +
99  c) * shapen_ + n) * shapey_ + y, j);
100  }
101 
102  private:
104  const index_t shapey_, shapez_, shapec_, shapen_;
105 };
106 template<typename SrcExp, typename DType, int dimsrc, int a2>
107 struct Plan<SwapAxisExp<SrcExp, DType, dimsrc, 1, a2>, DType> {
108  public:
110  : src_(MakePlan(e.src_)),
111  shapex_(e.shape_[dimsrc - 1]),
112  shapey_(e.shape_.ProdShape(a2 + 1, dimsrc - 1)),
113  shapez_(e.shape_[a2]) {}
114  MSHADOW_XINLINE DType Eval(index_t i, index_t x) const {
115  // swap x and z
116  const index_t y = i % shapey_;
117  i /= shapey_;
118  const index_t z = i % shapez_;
119  const index_t n = i / shapez_;
120  return src_.Eval((n * shapex_ + x) * shapey_ + y , z);
121  }
122 
123  private:
125  const index_t shapex_, shapey_, shapez_;
126 };
127 } // namespace expr
128 } // namespace mshadow
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