[Relay][Tensorflow] Allow an op as loop var.

  Allow binding a CallNode to a var to support op as loop var.

Author: Li Xiaoquan

include/tvm/relay/expr_functor.h

@@ -238,15 +238,15 @@ class ExprMutator
 void PostOrderVisit(const Expr& node, std::function<void(const Expr&)> fvisit);
 
 /*
- * \brief Bind function parameters or free variables.
+ * \brief Bind function parameters, free variables or CallNode.
  *
  * Parameter binding can only happen if expr is a Function.
  * binds cannot change internal arguments of internal functions.
  *
- * \param expr The function to be binded.
+ * \param expr The function to be bound.
  * \param binds The map of arguments to
  */
-Expr Bind(const Expr& expr, const tvm::Map<Var, Expr>& binds);
+Expr Bind(const Expr& expr, const tvm::Map<Expr, Expr>& binds);
 
 }  // namespace relay
 }  // namespace tvm

python/tvm/relay/frontend/tensorflow.py

@@ -1605,9 +1605,12 @@ def _while_loop(self):
         loop_vars = []
         bind_map = {}
         for i, var in enumerate(self.loop_vars):
-            assert isinstance(var, _expr.Var), repr(var)
-            v = tvm.relay.var("loop_var" + str(i),
-                              type_annotation=var.type_annotation)
+            if not isinstance(var, _expr.Var):
+                var_type = ir_pass.infer_type(var).checked_type
+            else:
+                var_type = var.type_annotation
+
+            v = tvm.relay.var("loop_var" + str(i), type_annotation=var_type)
             loop_vars.append(v)
             bind_map[var] = v
 

src/relay/ir/expr_functor.cc

@@ -355,7 +355,7 @@ TVM_REGISTER_API("relay._ir_pass.post_order_visit")
 // Implement bind.
 class ExprBinder : public ExprMutator {
  public:
-  explicit ExprBinder(const tvm::Map<Var, Expr>& args_map)
+  explicit ExprBinder(const tvm::Map<Expr, Expr>& args_map)
     : args_map_(args_map) {
   }
 
@@ -383,11 +383,21 @@ class ExprBinder : public ExprMutator {
     }
   }
 
+  Expr VisitExpr_(const CallNode* op) final {
+    auto id = GetRef<Call>(op);
+    auto it = args_map_.find(id);
+    if (it != args_map_.end()) {
+      return (*it).second;
+    } else {
+      return ExprMutator::VisitExpr_(op);
+    }
+  }
+
  private:
-  const tvm::Map<Var, Expr>& args_map_;
+  const tvm::Map<Expr, Expr>& args_map_;
 };
 
-Expr Bind(const Expr& expr, const tvm::Map<Var, Expr>& args_map) {
+Expr Bind(const Expr& expr, const tvm::Map<Expr, Expr>& args_map) {
   if (const FunctionNode* func = expr.as<FunctionNode>()) {
     Expr new_body = ExprBinder(args_map).Mutate(func->body);
     Array<Var> new_params;

tests/python/frontend/tensorflow/test_control_flow.py

@@ -51,6 +51,23 @@ def b(i): return tf.add(i, 1)
         check_equal(graph, tf_out)
 
 
+def test_callnode_loop_vars():
+    graph = tf.Graph()
+    with graph.as_default():
+        i = tf.add(tf.constant(0), 1)
+
+        def c(i): return tf.less(i, 10)
+
+        def b(i): return tf.add(i, 1)
+
+        r = tf.while_loop(c, b, [i])
+
+        with tf.Session() as sess:
+            tf_out = sess.run(r)
+
+        check_equal(graph, tf_out)
+
+
 def test_loop_2_vars():
     graph = tf.Graph()
     with graph.as_default():