[GitHub] [incubator-tvm] altanh commented on a change in pull request #6827: [RELAY][GRAD] Fix first-order AD on tuple arguments

[GitBox <gi...@apache.org>]

altanh commented on a change in pull request #6827:
URL: https://github.com/apache/incubator-tvm/pull/6827#discussion_r516360197



##########
File path: src/relay/transforms/gradient.cc
##########
@@ -181,6 +181,22 @@ struct FirstOrderReverseAD : ExprFunctor<ADValue(const Expr&)> {
     return ret;
   }
 
+  Expr UpdateGrad(const Type& t, const Expr& arg, const Expr& grad, LetList* ll) {
+    if (t.as<TensorTypeNode>()) {
+      return ll->Push(Add(arg, grad));
+    } else if (auto* tt = t.as<TupleTypeNode>()) {
+      Array<Expr> updates;
+      for (size_t i = 0; i < tt->fields.size(); ++i) {
+        updates.push_back(this->UpdateGrad(tt->fields[i], ll->Push(GetField(arg, i)),
+                                           ll->Push(GetField(grad, i)), ll));
+      }
+      return ll->Push(Tuple(updates));
+    } else {
+      LOG(FATAL) << "unsupported arg type of operator: " << t;

Review comment:
       I agree on this but it will probably need some refactoring, we might as well do it for the whole pass (first-order and higher-order). I think a separate PR will be ideal.




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org