Test cases updated

Author: siju-samuel

tests/python/relay/test_op_level5.py

@@ -45,49 +45,33 @@ def test_multibox_prior():
 
 
 def test_yolo_reorg():
-    ib = relay.ir_builder.IRBuilder()
     n, c, h, w = tvm.var("n"), tvm.var("c"), tvm.var("h"), tvm.var("w")
-    x = ib.param("x", relay.ty.TensorType((n, c, h, w), "float32"))
-    with ib.function(x) as func:
-        ib.ret(relay.vision.yolo_reorg(x))
-    ib.ret(func)
-    func = relay.ir_pass.infer_type(ib.env, func.to_func())
-    ftype = func.checked_type
-    assert ftype.ret_type == relay.ty.TensorType((n, c, h, w), "float32")
-
-    ib = relay.ir_builder.IRBuilder()
-    x = ib.param("x", relay.ty.TensorType((n, c, h, w), "float32"))
+    x = relay.var("x", relay.TensorType((n, c, h, w), "float32"))
+    z = relay.vision.yolo_reorg(x)
+    zz = relay.ir_pass.infer_type(z)
+    assert zz.checked_type == relay.ty.TensorType((n, c, h, w), "float32")
 
-    with ib.function(x) as func:
-        ib.ret(relay.vision.yolo_reorg(x, stride=2))
-    ib.ret(func)
-    func = relay.ir_pass.infer_type(ib.env, func.to_func())
-    ftype = func.checked_type
-    assert ftype.ret_type == relay.ty.TensorType((n, c*2*2, h/2, w/2), "float32")
+    x = relay.var("x", relay.TensorType((n, c, h, w), "float32"))
+    z = relay.vision.yolo_reorg(x, stride=2)
+    assert "stride=2" in z.astext()
+    zz = relay.ir_pass.infer_type(z)
+    assert zz.checked_type == relay.ty.TensorType((n, c*2*2, h/2, w/2), "float32")
 
 
 def test_yolo_region():
-    ib = relay.ir_builder.IRBuilder()
     n, c, h, w = tvm.var("n"), tvm.var("c"), tvm.var("h"), tvm.var("w")
-    x = ib.param("x", relay.ty.TensorType((n, c, h, w), "float32"))
-    with ib.function(x) as func:
-        ib.ret(relay.vision.yolo_region(x))
-    ib.ret(func)
-    func = relay.ir_pass.infer_type(ib.env, func.to_func())
-    ftype = func.checked_type
-    assert ftype.ret_type == relay.ty.TensorType((n, c, h, w), "float32")
+    x = relay.var("x", relay.TensorType((n, c, h, w), "float32"))
+    z = relay.vision.yolo_region(x)
+    zz = relay.ir_pass.infer_type(z)
+    assert zz.checked_type == relay.ty.TensorType((n, c, h, w), "float32")
 
 
 def test_yolov3_yolo():
-    ib = relay.ir_builder.IRBuilder()
     n, c, h, w = tvm.var("n"), tvm.var("c"), tvm.var("h"), tvm.var("w")
-    x = ib.param("x", relay.ty.TensorType((n, c, h, w), "float32"))
-    with ib.function(x) as func:
-        ib.ret(relay.vision.yolov3_yolo(x))
-    ib.ret(func)
-    func = relay.ir_pass.infer_type(ib.env, func.to_func())
-    ftype = func.checked_type
-    assert ftype.ret_type == relay.ty.TensorType((n, c, h, w), "float32")
+    x = relay.var("x", relay.TensorType((n, c, h, w), "float32"))
+    z = relay.vision.yolov3_yolo(x)
+    zz = relay.ir_pass.infer_type(z)
+    assert zz.checked_type == relay.ty.TensorType((n, c, h, w), "float32")
 
 
 if __name__ == "__main__":