fix lint

tlc-pack · Apr 4, 2022 · 75c61da · 75c61da
1 parent fd423ef
commit 75c61da
Show file tree

Hide file tree

Showing 4 changed files with 4 additions and 9 deletions.
diff --git a/python/tvm/script/relax/parser.py b/python/tvm/script/relax/parser.py
@@ -1332,7 +1332,7 @@ def transform_block(self, block: ast.Block) -> relax.SeqExpr:
                     blocks.append(relax.BindingBlock(current_block, self.to_tvm_span(stmt.span)))
                     current_block = []
                 blocks.append(parsed_stmt)
-            elif isinstance(parsed_stmt, relax.Function) or isinstance(parsed_stmt, tir.PrimFunc):
+            elif isinstance(parsed_stmt, (relax.Function, tir.PrimFunc)):
                 func_var = self.decl_var(stmt.name, None, None, stmt.span)
                 current_block.append(
                     relax.VarBinding(func_var, parsed_stmt, self.to_tvm_span(stmt.span))

diff --git a/src/relax/analysis/analysis.cc b/src/relax/analysis/analysis.cc
@@ -156,12 +156,7 @@ tvm::Array<GlobalVar> RecGlobalVars(const Expr& expr) { return VarVisitor().RecG
 
 TVM_REGISTER_GLOBAL("relax.analysis.free_vars").set_body_typed(FreeVars);
 
-TVM_REGISTER_GLOBAL("relax.analysis.bound_vars").set_body([](TVMArgs args, TVMRetValue* ret) {
-  ObjectRef x = args[0];
-  if (x.as<ExprNode>()) {
-    *ret = BoundVars(Downcast<Expr>(x));
-  }
-});
+TVM_REGISTER_GLOBAL("relax.analysis.bound_vars").set_body_typed(BoundVars);
 
 TVM_REGISTER_GLOBAL("relax.analysis.all_vars").set_body_typed(AllVars);
 

diff --git a/src/relax/transform/lambda_lift.cc b/src/relax/transform/lambda_lift.cc
@@ -18,7 +18,7 @@
  */
 
 /*!
- * \file tvm/relax/backend/vm/lambda_lift.cc
+ * \file tvm/relax/transform/lambda_lift.cc
  * \brief Lift local functions into global functions.
  */
 

diff --git a/tests/python/relax/test_lambda_lift.py b/tests/python/relax/test_lambda_lift.py
@@ -175,7 +175,7 @@ def while_loop(i: Tensor[(), "int32"], s: Tensor[(2, 3), "float32"]):
     # Perform Lamda Lifting
     after = transform.LambdaLift()(before)
     assert len(after.functions) == 2
-    assert_structural_equal(after["lifted_func_0"], expected["lifted_func_0"], map_free_vars=True)
+    assert_structural_equal(after, expected, map_free_vars=True)
     _check_save_roundtrip(after)
-Original file line number
+Diff line change
@@ Expand Up / @@ -18,7 +18,7 @@ @@
      */
     /*!
-     * \file tvm/relax/backend/vm/lambda_lift.cc
+     * \file tvm/relax/transform/lambda_lift.cc
      * \brief Lift local functions into global functions.
      */
@@ Expand Down @@