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

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

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



##########
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:
       Can we try and do diagnostics here? we could put into improve AD with diagnostics 




----------------------------------------------------------------
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