def test_recursion(): """ Program: let f(n: i32, data: f32) -> f32 = { if (n == 0) { return data; } else { return f(n - 1, log(data)); } } f(2, 10000); """ f = relay.Var("f") n = relay.Var("n") np = relay.Param(n, e.int32) data = relay.Var("data") datap = relay.Param(data, e.float32) funcbody = relay.If(equal(n, convert(0)), data, f(subtract(n, convert(1.0)), log(data))) value = relay.Function([np, datap], e.float32, funcbody, []) orig = relay.Let(f, funcbody, f(convert(2.0), convert(10000.0)), e.float32) assert alpha_equal(dead_code_elimination(orig), orig) assert alpha_equal( dead_code_elimination(relay.Let(f, funcbody, e.three, e.float32)), e.three)
def test_param_alpha_equal(): # only checks equality of the types v1 = relay.Var("v1") v2 = relay.Var("v2") p1 = relay.Param(v1, relay.TensorType((1, 2, 3), "float32")) p2 = relay.Param(v2, relay.TensorType((1, 2, 3), "float32")) assert alpha_equal(p1, p2) p3 = relay.Param(v1, relay.TensorType((4, 5, 6), "int8")) assert not alpha_equal(p1, p3) p4 = relay.Param(v1, relay.TupleType([relay.TensorType((1, 2, 3), "float32")])) assert not alpha_equal(p1, p4)
def test_param(): lv = relay.Var('x') ty = None param = relay.Param(lv, ty) assert param.var == lv assert param.type == ty assert param.span == None str(param)
def test_function(): param_names = ['a', 'b', 'c', 'd'] params = tvm.convert([relay.Param(relay.Var(n), None) for n in param_names]) ret_type = None body = params[0].var type_params = tvm.convert([]) fn = relay.Function(params, ret_type, body, type_params) show(fn)
def test_function(): param_names = ['a', 'b', 'c', 'd'] params = tvm.convert([relay.Param(relay.Var(n), None) for n in param_names]) ret_type = None body = None type_params = tvm.convert([]) fn = relay.Function(params, ret_type, body, type_params) assert fn.params == params assert fn.body == body assert fn.type_params == type_params assert fn.span == None str(fn)
def test_free_vars(): x = relay.Var("x") fvx = free_vars(x) assert len(fvx) == 1 assert fvx[0] == x v = relay.Constant(tvm.nd.array(10)) ty = relay.TensorType([], "int32") let = relay.Let(x, v, x, ty) fvx = free_vars(let) assert len(free_vars(let)) == 0 f = relay.Function([relay.Param(x, ty)], ty, x) assert len(free_vars(f)) == 0
def test_well_formed(): x = relay.Var('x') assert well_formed(x) v = relay.Constant(tvm.nd.array(10)) ty = None let = relay.Let(x, v, x, ty) assert well_formed(let) assert not well_formed(relay.Let(x, v, let, ty)) f = relay.Function([relay.Param(x, ty)], ty, x) assert well_formed(f) # this test should pass in case of weak uniqueness (only test for shadowing) # but we want all binder to be distinct from each other. assert not well_formed(relay.Let(relay.Var("y"), f, relay.Let(relay.Var("z"), f, v, ty), ty))
def test_param(): lv = relay.Var('x') ty = None param = relay.Param(lv, ty) show(lv)
def test_function_alpha_equal(): v1 = relay.Var("v1") v2 = relay.Var("v2") v3 = relay.Var("v3") v4 = relay.Var("v4") tt1 = relay.TensorType((1, 2, 3), "float32") tt2 = relay.TensorType((4, 5, 6), "int8") tt3 = relay.TupleType([tt1, tt2]) tp1 = relay.TypeParam("tp1", relay.Kind.Type) tp2 = relay.TypeParam("tp2", relay.Kind.Type) tp3 = relay.TypeParam("tp3", relay.Kind.Shape) tp4 = relay.TypeParam("tp4", relay.Kind.Shape) basic_args = [relay.Param(v3, tt1), relay.Param(v4, tt2)] basic_tps = [tp1, tp2] func = relay.Function([relay.Param(v1, tt1), relay.Param(v2, tt2)], tt2, v2, basic_tps) mapped = relay.Function(basic_args, tt2, v4, basic_tps) assert alpha_equal(func, mapped) fewer_params = relay.Function([relay.Param(v4, tt2)], tt2, v4, basic_tps) assert not alpha_equal(func, fewer_params) more_params = relay.Function([relay.Param(v3, tt1), relay.Param(v4, tt2), relay.Param(v2, tt2)], tt2, v4, basic_tps) assert not alpha_equal(func, more_params) params_unordered = relay.Function([relay.Param(v3, tt2), relay.Param(v4, tt1)], tt1, v3, basic_tps) assert not alpha_equal(func, params_unordered) params_mismatch = relay.Function([relay.Param(v3, tt3), relay.Param(v4, tt2)], tt2, v4, basic_tps) assert not alpha_equal(func, params_mismatch) # also would not typecheck ret_type_mismatch = relay.Function(basic_args, tt1, v4, basic_tps) assert not alpha_equal(func, ret_type_mismatch) # also mis-typed different_body = relay.Function(basic_args, tt2, v3, basic_tps) assert not alpha_equal(func, different_body) fewer_type_params = relay.Function(basic_args, tt2, v4, [tp1]) assert not alpha_equal(func, fewer_type_params) more_type_params = relay.Function(basic_args, tt2, v4, [tp1, tp2, tp3]) assert not alpha_equal(func, more_type_params) type_params_unordered = relay.Function(basic_args, tt2, v4, [tp2, tp1]) assert not alpha_equal(func, type_params_unordered) different_type_params = relay.Function(basic_args, tt2, v4, [tp3, tp4]) assert not alpha_equal(func, different_type_params) # a well-typed example that also differs in body, ret type, and type params tupled_example = relay.Function(basic_args, tt3, relay.Tuple([v3, v4])) assert not alpha_equal(func, tupled_example)