MCPcopy Create free account
hub / github.com/MegEngine/MegEngine / optimize_grad_var

Method optimize_grad_var

src/opr/impl/loop/grad.cpp:143–167  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

141}
142
143SymbolVar OutputRecorderSumIntoDest::optimize_grad_var(SymbolVar grad) {
144 if (m_sum_last)
145 return grad;
146
147 // currently only try to coalesce incr_sub oprs
148 auto opr = grad.node()->owner_opr();
149 if (!gopt::check_is_incr_subtensor_zero(opr))
150 return grad;
151
152 // now we are sure that grad is in the form of incr_sub(0, sub)
153 m_optimize_coalesce_incr_sub = true;
154 test_check_optimize_success = true;
155
156 {
157 using namespace std::placeholders;
158 auto replacer = std::bind(
159 &OutputRecorderSumIntoDest::incr_sub_input_replacer, this, _1);
160 m_grad_incr_subtensor_modifier =
161 gopt::remake_incr_subtensor_zero(opr, nullptr, replacer);
162 }
163
164 // use a placeholder grad var to ensure sub is computed; result is computed
165 // correctly since m_optimize_coalesce_incr_sub has been set
166 return opr->input(1);
167}
168
169DeviceTensorND OutputRecorderSumIntoDest::incr_sub_input_replacer(
170 const TensorShape& shape) {

Callers 1

bind_grad_varMethod · 0.80

Calls 3

owner_oprMethod · 0.45
nodeMethod · 0.45
inputMethod · 0.45

Tested by

no test coverage detected