mxnet
take_grad.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_TAKE_GRAD_H_
27 #define MSHADOW_EXTENSION_TAKE_GRAD_H_
28 
29 #include "../extension.h"
30 
31 namespace mshadow {
32 namespace expr {
33 
40 template<typename IndexExp, typename SrcExp, typename DType>
41 struct TakeGradExp : public Exp<TakeGradExp<IndexExp, SrcExp, DType>,
42  DType, type::kChainer> {
44  const IndexExp &index_;
46  const SrcExp &src_;
50  TakeGradExp(const IndexExp &index, const SrcExp &src, const index_t input_dim)
51  : index_(index), src_(src), input_dim_(input_dim) {}
52 }; // struct TakeGradExp
53 
54 
55 template<typename IndexExp,
56  typename SrcExp,
57  typename DType,
58  int e1, int e2>
61  const Exp<SrcExp, DType, e2> &src,
62  const index_t input_dim) {
64  src.self(),
65  input_dim);
66 }
67 
68 //----------------------
69 // Execution plan
70 //----------------------
71 
72 template<typename IndexExp, typename SrcExp, typename DType>
73 struct Plan<TakeGradExp<IndexExp, SrcExp, DType>, DType> {
74  public:
76  : index_(MakePlan(e.index_)),
77  src_(MakePlan(e.src_)),
78  batch_size_(ShapeCheck<1, IndexExp>::Check(e.index_)[0]) {
79  }
80 
81  // now return shape: in * out
82  MSHADOW_XINLINE DType Eval(index_t y, index_t x) const {
83  DType ret = 0.f;
84  for (index_t i = 0; i < batch_size_; ++i) {
85  index_t idx = static_cast<index_t>(index_.Eval(0, i));
86  if (idx == y) {
87  ret += static_cast<DType>(src_.Eval(i, x));
88  }
89  }
90  return ret;
91  }
92 
93  private:
96  const index_t batch_size_;
97 }; // struct Plan
98 
99 
100 template<typename IndexExp, typename SrcExp, typename DType>
103  return Plan<TakeGradExp<IndexExp, SrcExp, DType>, DType>(exp);
104 }
105 
106 template<int dim, typename IndexExp, typename SrcExp, typename DType>
107 struct ShapeCheck<dim, TakeGradExp<IndexExp, SrcExp, DType> > {
108  inline static Shape<dim>
110  CHECK(dim == 2)
111  << "TakeGradExp only support 2D output";
112  // Shape<1> dshape = ShapeCheck<1, IndexExp>::Check(t.index_);
114  Shape<dim> ret;
115  ret[0] = t.input_dim_;
116  ret[1] = gshape[1];
117  return ret;
118  }
119 }; // struct ShapeCheck
120 
121 template<typename IndexExp, typename SrcExp, typename DType>
122 struct ExpInfo<TakeGradExp<IndexExp, SrcExp, DType> > {
123  static const int kDim = 2;
124  static const int kDevMask = ExpInfo<IndexExp>::kDevMask;
125 };
126 
127 } // namespace expr
128 } // namespace mshadow
129 
130 #endif // MSHADOW_EXTENSION_TAKE_GRAD_H_
const IndexExp & index_
index oprand
Definition: take_grad.h:44
Definition: expr_engine-inl.h:59
const SrcExp & src_
out gradient oprand
Definition: take_grad.h:46
static Shape< dim > Check(const E &t)
#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
static Shape< dim > Check(const TakeGradExp< IndexExp, SrcExp, DType > &t)
Definition: take_grad.h:109
int32_t index_t
type that will be used for index
Definition: base.h:336
TakeGradExp< IndexExp, SrcExp, DType > take_grad(const Exp< IndexExp, DType, e1 > &index, const Exp< SrcExp, DType, e2 > &src, const index_t input_dim)
Definition: take_grad.h:60
runtime shape checking template get the shape of an expression, report error if shape mismatch ...
Definition: expr_engine-inl.h:365
Calculate embedding gradient.
Definition: take_grad.h:41
const index_t input_dim_
batch size
Definition: take_grad.h:48
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
TakeGradExp(const IndexExp &index, const SrcExp &src, const index_t input_dim)
constructor
Definition: take_grad.h:50
MSHADOW_XINLINE DType Eval(index_t y, index_t x) const
Definition: take_grad.h:82
overloaded + operator between half_t and bf16_t
Definition: base.h:327
Plan(const TakeGradExp< IndexExp, SrcExp, DType > &e)
Definition: take_grad.h:75