fix error

src/relay/ir/adt.cc

@@ -6,9 +6,9 @@
  * to you under the Apache License, Version 2.0 (the
  * "License"); you may not use this file except in compliance
  * with the License.  You may obtain a copy of the License at
- * 
+ *
  *   http://www.apache.org/licenses/LICENSE-2.0
- * 
+ *
  * Unless required by applicable law or agreed to in writing,
  * software distributed under the License is distributed on an
  * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
@@ -18,7 +18,7 @@
  */
 
 /*!
- *  Copyright (c) 2018 by Contributors
+ *  Copyright (c) 2019 by Contributors
  * \file src/tvm/ir/adt.cc
  * \brief AST nodes for Relay algebraic data types (ADTs).
  */

src/relay/pass/dependency_graph.h

@@ -20,7 +20,7 @@
 /*!
  *  Copyright (c) 2019 by Contributors.
  * \file tvm/relay/pass/dependency_graph.h
- * \brief 
+ * \brief create a dependency graph.
  */
 #ifndef TVM_RELAY_PASS_DEPENDENCY_GRAPH_H_
 #define TVM_RELAY_PASS_DEPENDENCY_GRAPH_H_

src/relay/pass/let_list.h

@@ -6,9 +6,9 @@
  * to you under the Apache License, Version 2.0 (the
  * "License"); you may not use this file except in compliance
  * with the License.  You may obtain a copy of the License at
- * 
+ *
  *   http://www.apache.org/licenses/LICENSE-2.0
- * 
+ *
  * Unless required by applicable law or agreed to in writing,
  * software distributed under the License is distributed on an
  * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
@@ -18,7 +18,7 @@
  */
 
 /*!
- *  Copyright (c) 2018 by Contributors
+ *  Copyright (c) 2019 by Contributors
  * \file let_list.h
  * \brief LetList record let binding and insert let expression implicitly.
  *  using it, one can treat AST as value instead of expression,
@@ -46,6 +46,11 @@ namespace relay {
  */
 class LetList {
  public:
+  ~LetList() {
+    if (lets_.size() > 0 && !used_) {
+      std::cout << "Warning: letlist not used" << std::endl;
+    }
+  }
   /*!
    * \brief insert a binding.
    *

src/relay/pass/to_a_normal_form.cc

@@ -72,13 +72,16 @@ Scope LCA(Scope lhs, Scope rhs) {
 
 std::unordered_map<DependencyGraph::Node*, Scope> CalcScope(const DependencyGraph& dg) {
   std::unordered_map<DependencyGraph::Node*, Scope> expr_scope;
+  bool global_scope_used = false;
   Scope global_scope = std::make_shared<ScopeNode>();
   for (auto it = dg.post_dfs_order.rbegin(); it != dg.post_dfs_order.rend(); ++it) {
     DependencyGraph::Node* n = *it;
     auto iit = n->parents.head;
     Scope s;
     if (iit == nullptr) {
+      CHECK(!global_scope_used);
       s = global_scope;
+      global_scope_used = true;
     } else {
       s = expr_scope.at(iit->value);
       iit = iit->next;
@@ -88,6 +91,7 @@ std::unordered_map<DependencyGraph::Node*, Scope> CalcScope(const DependencyGrap
     }
     expr_scope.insert({n, n->new_scope ? ChildScope(s) : s});
   }
+  CHECK(global_scope_used);
   return expr_scope;
 }
 
@@ -133,22 +137,26 @@ class Fill : ExprFunctor<Expr(const Expr&, const Var&)> {
   Expr VisitExpr(const Expr& e, const Var& v) final {
     if (memo.count(e) == 0) {
       memo.insert({e, ExprFunctor<Expr(const Expr&, const Var&)>::VisitExpr(e, v)});
+    } else if (v.defined()) {
+      GetScope(e)->ll->Push(v, memo.at(e));
     }
-    return memo.at(e);
+    auto ret = memo.at(e);
+    CHECK(IsAtomic(ret));
+    return ret;
   }
 
   Expr VisitExpr(const Expr& e) {
     return this->VisitExpr(e, Var());
   }
 
-  Expr Atomic(const Expr& orig, const Expr& now, const Var& v) {
-    return v.defined() ? GetScope(orig)->ll->Push(v, now) : now;
+  Expr Atomic(const Expr& e, const Var& v) {
+    return v.defined() ? GetScope(e)->ll->Push(v, e) : e;
   }
 
   Expr Compound(const Expr& orig, const Expr& now, const Var& v) {
     Var var = v.defined() ?
       v :
-      VarNode::make(std::string("x"), IncompleteTypeNode::make(Kind::kType));
+      VarNode::make(std::string("x"), Type());
     return GetScope(orig)->ll->Push(var, now);
   }
 
@@ -227,22 +235,22 @@ class Fill : ExprFunctor<Expr(const Expr&, const Var&)> {
 
   Expr VisitExpr_(const VarNode* vn, const Var& v) final {
     Expr e = GetRef<Expr>(vn);
-    return Atomic(e, e, v);
+    return Atomic(e, v);
   }
 
   Expr VisitExpr_(const GlobalVarNode* gvn, const Var& v) final {
     GlobalVar gv = GetRef<GlobalVar>(gvn);
-    return Atomic(gv, gv, v);
+    return Atomic(gv, v);
   }
 
   Expr VisitExpr_(const OpNode* op, const Var& v) final {
     Expr e = GetRef<Expr>(op);
-    return Atomic(e, e, v);
+    return Atomic(e, v);
   }
 
   Expr VisitExpr_(const ConstructorNode* c, const Var& v) final {
     Expr e = GetRef<Expr>(c);
-    return Atomic(e, e, v);
+    return Atomic(e, v);
   }
 
   Expr VisitExpr_(const MatchNode* m, const Var& v) final {
@@ -290,11 +298,12 @@ Module ToANormalForm(const Module& m) {
   tvm::Map<GlobalVar, Function> updates;
   auto funcs = m->functions;
   for (const auto& it : funcs) {
+    CHECK_EQ(FreeVars(it.second).size(), 0);
     Expr ret =
       TransformF([&](const Expr& e) {
         return ToANormalFormAux(e);
       }, it.second);
-    CHECK_EQ(FreeVars(ret).size(), 0);
+    CHECK_EQ(FreeVars(ret).size(), 0) << AsText(ret) << "should not has free vars: " << FreeVars(ret);
     updates.Set(it.first, Downcast<Function>(ret));
   }
 

tests/python/relay/test_pass_to_a_normal_form.py

@@ -22,6 +22,7 @@
 from tvm.relay.prelude import Prelude
 from tvm.relay.testing import add_nat_definitions, count
 from tvm.relay.feature import Feature
+import tvm.relay.testing
 
 
 def check_eval(expr, expected_result, mod=None, rtol=1e-07):
@@ -180,6 +181,20 @@ def test_function():
     check_eval(anf_f(d), 8)
 
 
+def test_gradient_if():
+    x = relay.var("a", shape=(1, 16))
+    y = relay.var("y", shape=(1, 16))
+    cond = relay.var("cond", shape=(), dtype='uint1')
+    net = relay.If(cond, x, x)
+    net = relay.add(x, net)
+    net = relay.Function([cond,x,y], net)
+    net = relay.ir_pass.infer_type(net)
+    mod = relay.Module.from_expr(net)
+    mod = relay.transform.ToANormalForm()(mod)
+    mod[mod.entry_func] = relay.ir_pass.gradient(mod[mod.entry_func], mode='higher_order')
+    mod = relay.transform.ToANormalForm()(mod)
+
+
 if __name__ == '__main__':
     test_explicit_bound()
     test_order()
@@ -189,3 +204,4 @@ def test_function():
     test_let()
     test_nat_add()
     test_function()
+    test_gradient_if()