[Bugfix] Simplify reduce expression in te.gradient (#6611)

apache · Oct 2, 2020 · 1ff5f39 · 1ff5f39
1 parent 9f5b9da
commit 1ff5f39
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/te/autodiff/ad_simplify.cc b/src/te/autodiff/ad_simplify.cc
@@ -1172,8 +1172,10 @@ PrimExpr RemoveJacobianAndLiftNonzeroCondImpl(const PrimExpr& expr_orig, const A
 
       new_red = Reduce(red->combiner, source, red->axis, cond, red->value_index, red->init);
       new_red = SimplifyReductionDomain(new_red, combined_vranges);
+      // Update original red pointer for later use.
+      red = new_red.as<ReduceNode>();
       // If the reduction disappears completely then transform the result as a non-reduction
-      if (!new_red.as<ReduceNode>()) {
+      if (!red) {
         return RemoveJacobianAndLiftNonzeroCondImpl(new_red, axis, vranges);
       }