def test_free_type_vars(): tp = relay.TypeVar("") ty = relay.TupleType([tp, relay.TensorType([], "int32")]) x = relay.Var("x", ty) y = relay.Var("y") let = relay.Let(x, y, x) fvl = free_vars(let) assert len(fvl) == 1 assert fvl[0] == y ftvl = free_type_vars(let) assert len(ftvl) == 1 assert ftvl[0] == tp
def test_rev(): a = relay.TypeVar("a") assert mod[rev].checked_type == relay.FuncType([l(a)], l(a), [a]) res = intrp.evaluate( rev(cons(build_nat(1), cons(build_nat(2), cons(build_nat(3), nil()))))) reversed = to_list(res) assert len(reversed) == 3 assert count(reversed[0]) == 3 assert count(reversed[1]) == 2 assert count(reversed[2]) == 1
def test_extern_adt_defn(): # TODO(weberlo): update this test once extern is implemented mod = tvm.IRModule() extern_var = relay.GlobalTypeVar("T") typ_var = relay.TypeVar("A") extern_def = relay.TypeData(extern_var, [typ_var], []) mod[extern_var] = extern_def assert_parses_as(""" extern type T[A] """, mod)
def test_foldl(): a = relay.TypeVar("a") b = relay.TypeVar("b") lhs = mod[foldl].checked_type rhs = relay.FuncType([relay.FuncType([a, b], a), a, l(b)], a, [a, b]) assert lhs == rhs x = relay.Var("x") y = relay.Var("y") rev_dup = relay.Function([y, x], cons(x, cons(x, y))) res = intrp.evaluate( foldl( rev_dup, nil(), cons(make_nat_expr(1), cons(make_nat_expr(2), cons(make_nat_expr(3), nil()))), ) ) reversed = to_list(res) assert len(reversed) == 6 assert count(reversed[0]) == 3 and count(reversed[1]) == 3 assert count(reversed[2]) == 2 and count(reversed[3]) == 2 assert count(reversed[4]) == 1 and count(reversed[5]) == 1
def test_type_relation(): tp = relay.TypeVar('tp', relay.Kind.Type) tf = relay.FuncType(tvm.convert([]), None, tvm.convert([]), tvm.convert([])) tt = relay.TensorType(tvm.convert([1, 2, 3]), 'float32') args = tvm.convert([tf, tt, tp]) num_inputs = 2 func = tvm.get_env_func("tvm.relay.type_relation.Broadcast") attrs = tvm.make.node("attrs.TestAttrs", name="attr", padding=(3,4)) tr = relay.TypeRelation(func, args, num_inputs, attrs) assert tr.args == args assert tr.num_inputs == num_inputs
def test_zip(): a = relay.TypeVar("a") b = relay.TypeVar("b") expected_type = relay.FuncType([l(a), l(b)], l(relay.TupleType([a, b])), [a, b]) assert mod[zip].checked_type == expected_type l1 = cons(make_nat_expr(1), cons(make_nat_expr(2), cons(make_nat_expr(3), nil()))) l2 = cons( nil(), cons(cons(nil(), nil()), cons(cons(nil(), cons(nil(), nil())), nil()))) res = intrp.evaluate(zip(l1, l2)) zipped = to_list(res) assert len(zipped) == 3 assert count(zipped[0][0]) == 1 assert len(to_list(zipped[0][1])) == 0 assert count(zipped[1][0]) == 2 assert len(to_list(zipped[1][1])) == 1 assert count(zipped[2][0]) == 3 assert len(to_list(zipped[2][1])) == 2 # test truncation l3 = cons(make_nat_expr(4), cons(make_nat_expr(5), nil())) shorter_res = intrp.evaluate(zip(l3, l2)) truncated = to_list(shorter_res) assert len(truncated) == 2 assert count(truncated[0][0]) == 4 assert len(to_list(truncated[0][1])) == 0 assert count(truncated[1][0]) == 5 assert len(to_list(truncated[1][1])) == 1 l4 = cons(nil(), nil()) shortest_res = intrp.evaluate(zip(l3, l4)) singleton = to_list(shortest_res) assert len(singleton) == 1 assert count(singleton[0][0]) == 4 assert len(to_list(singleton[0][1])) == 0
def test_constructor_type(): mod = tvm.IRModule() box, constructor = initialize_box_adt(mod) a = relay.TypeVar("a") x = relay.Var("x", a) func = relay.Function([x], constructor(x), box(a), [a]) mod["main"] = func mod = infer_mod(mod) func_ty = mod["main"].checked_type box = mod.get_global_type_var("box") expected = relay.FuncType([a], box(a), [a]) assert func_ty == expected
def test_multiple_type_param_defn(): glob_typ_var = relay.GlobalTypeVar("Either") typ_var_a = relay.TypeVar("A") typ_var_b = relay.TypeVar("B") prog = relay.TypeData( glob_typ_var, [typ_var_a, typ_var_b], [ relay.Constructor("Left", [typ_var_a], glob_typ_var), relay.Constructor("Right", [typ_var_b], glob_typ_var), ]) mod = relay.Module() mod[glob_typ_var] = prog assert parses_as( """ type Either[A, B] { Left(A), Right(B), } """, mod )
def test_foldr(): a = relay.TypeVar("a") b = relay.TypeVar("b") lhs = prelude.mod[foldr].checked_type rhs = relay.FuncType([relay.FuncType([a, b], b), b, rlist(a)], b, [a, b]) assert lhs == rhs x = relay.Var("x") y = relay.Var("y") identity = relay.Function([x, y], cons(x, y)) res = intrp.evaluate( foldr( identity, nil(), cons( make_nat_expr(prelude, 1), cons(make_nat_expr(prelude, 2), cons(make_nat_expr(prelude, 3), nil())), ), ) ) same = to_list(res) assert len(same) == 3 assert count(same[0]) == 1 and count(same[1]) == 2 and count(same[2]) == 3
def test_extern_adt_defn(): mod = tvm.IRModule() extern_var = relay.GlobalTypeVar("T") typ_var = relay.TypeVar("A") extern_def = relay.TypeData(extern_var, [typ_var], []) mod[extern_var] = extern_def assert_parse_module_as( """ extern type T[A] """, mod, )
def test_multiple_cons_defn(): mod = relay.Module() list_var = relay.GlobalTypeVar("List") typ_var = relay.TypeVar("A") prog = relay.TypeData( list_var, [typ_var], [ relay.Constructor("Cons", [typ_var, list_var(typ_var)], list_var), relay.Constructor("Nil", [], list_var), ]) mod[list_var] = prog assert parses_as(LIST_DEFN, mod)
def test_concat(): a = relay.TypeVar("a") assert mod[concat].checked_type == relay.FuncType([l(a), l(a)], l(a), [a]) l1 = cons(make_nat_expr(1), cons(make_nat_expr(2), nil())) l2 = cons(make_nat_expr(3), cons(make_nat_expr(4), nil())) res = intrp.evaluate(concat(l1, l2)) catted = to_list(res) assert len(catted) == 4 assert count(catted[0]) == 1 assert count(catted[1]) == 2 assert count(catted[2]) == 3 assert count(catted[3]) == 4
def test_type_param_sequal(): t1 = relay.TypeVar("v1", relay.TypeKind.Type) t2 = relay.TypeVar("v2", relay.TypeKind.ShapeVar) t3 = relay.TypeVar("v3", relay.TypeKind.Type) # only pointer equality and eq_map allow equal params assert t1 == t1 assert t2 == t2 assert t1 != t2 # different kind assert t1 != t3 # not in eq_map # function types are the only way to put type params # in eq map ft1 = relay.FuncType(tvm.runtime.convert([]), t1, tvm.runtime.convert([t1]), tvm.runtime.convert([])) ft2 = relay.FuncType(tvm.runtime.convert([]), t3, tvm.runtime.convert([t3]), tvm.runtime.convert([])) # actually an invalid type because t2 is wrong kind ft3 = relay.FuncType(tvm.runtime.convert([]), t2, tvm.runtime.convert([t2]), tvm.runtime.convert([])) assert ft1 == ft2 assert ft1 != ft3 # kinds still do not match
def test_foldr1(): a = relay.TypeVar("a") lhs = mod[p.foldr1].checked_type rhs = relay.FuncType([relay.FuncType([a, a], a), l(a)], a, [a]) assert lhs == rhs x = relay.Var("x") y = relay.Var("y") f = relay.Function([x, y], add(x, y)) res = intrp.evaluate( foldr1(f, cons(make_nat_expr(1), cons(make_nat_expr(2), cons(make_nat_expr(3), nil())))) ) assert count(res) == 6
def test_unfoldl(): a = relay.TypeVar("a") b = relay.TypeVar("b") expected_type = relay.FuncType( [relay.FuncType([a], optional(relay.TupleType([a, b]))), a], l(b), [a, b]) x = relay.Var("x", nat()) n = relay.Var("n", nat()) count_down = relay.Function( [x], relay.Match(x, [ relay.Clause(relay.PatternConstructor(s, [relay.PatternVar(n)]), some(relay.Tuple([n, x]))), relay.Clause(relay.PatternConstructor(z, []), none()) ])) res = intrp.evaluate(unfoldl(count_down, make_nat_expr(3))) unfolded = to_list(res) assert len(unfolded) == 3 assert count(unfolded[0]) == 1 assert count(unfolded[1]) == 2 assert count(unfolded[2]) == 3
def test_self_reference(): """ Program: def f(x) { return x; } """ a = relay.TypeVar("a") x = relay.var("x", a) sb = relay.ScopeBuilder() f = relay.Function([x], x) fx = relay.Call(f, [x]) assert relay.ir_pass.infer_type(x).checked_type == a assert relay.ir_pass.infer_type(f).checked_type == relay.FuncType([a], a) assert relay.ir_pass.infer_type(fx).checked_type == a
def test_type_relation(): tp = relay.TypeVar('tp', relay.TypeKind.Type) tf = relay.FuncType(tvm.convert([]), None, tvm.convert([]), tvm.convert([])) tt = relay.TensorType(tvm.convert([1, 2, 3]), 'float32') args = tvm.convert([tp, tf, tt]) num_inputs = 2 func = tvm.ir.EnvFunc.get("tvm.relay.type_relation.Broadcast") attrs = tvm.ir.make_node("attrs.TestAttrs", name="attr", padding=(3,4)) tr = relay.TypeRelation(func, args, num_inputs, attrs) assert tr.args == args assert tr.num_inputs == num_inputs str(tr) check_json_roundtrip(tr)
def test_typecall_kind(): gtv = relay.GlobalTypeVar("gtv") mod = tvm.IRModule() data = relay.TypeData(gtv, [], []) mod[gtv] = data empty_call = relay.TypeCall(gtv, []) assert check_kind(empty_call, mod) == relay.TypeKind.Type new_mod = tvm.IRModule() tv = relay.TypeVar("tv") new_data = relay.TypeData(gtv, [tv], []) new_mod[gtv] = new_data call = relay.TypeCall(gtv, [relay.TupleType([])]) assert check_kind(call, new_mod) == relay.TypeKind.Type
def test_map_accuml(): a = relay.TypeVar("a") b = relay.TypeVar("b") c = relay.TypeVar("c") expected_type = relay.FuncType( [relay.FuncType([a, b], relay.TupleType([a, c])), a, l(b)], relay.TupleType([a, l(c)]), [a, b, c]) assert mod[map_accuml].checked_type == expected_type acc = relay.Var("acc", nat()) x = relay.Var("x", nat()) add_to_acc = relay.Function([acc, x], relay.Tuple([add(x, acc), x])) vals = cons(build_nat(1), cons(build_nat(2), cons(build_nat(3), nil()))) res = intrp.evaluate(map_accuml(add_to_acc, z(), vals)) sum = count(res[0]) new_vals = to_list(res[1]) assert sum == 6 assert len(new_vals) == 3 assert count(new_vals[0]) == 3 assert count(new_vals[1]) == 2 assert count(new_vals[2]) == 1
def test_nested_matches(): a = relay.TypeVar("a") x = relay.Var("x") y = relay.Var("y") w = relay.Var("w") h = relay.Var("h") t = relay.Var("t") flatten = relay.GlobalVar("flatten") # flatten could be written using a fold, but this way has nested matches inner_match = relay.Match( y, [ relay.Clause(relay.PatternConstructor(nil), flatten(w)), relay.Clause( relay.PatternConstructor(cons, [relay.PatternVar(h), relay.PatternVar(t)]), cons(h, flatten(cons(t, w))), ), ], ) mod[flatten] = relay.Function( [x], relay.Match( x, [ relay.Clause(relay.PatternConstructor(nil), nil()), relay.Clause( relay.PatternConstructor(cons, [relay.PatternVar(y), relay.PatternVar(w)]), inner_match, ), ], ), l(a), [a], ) first_list = cons(make_nat_expr(1), cons(make_nat_expr(2), cons(make_nat_expr(3), nil()))) second_list = cons(make_nat_expr(4), cons(make_nat_expr(5), cons(make_nat_expr(6), nil()))) final_list = cons(first_list, cons(second_list, nil())) res = intrp.evaluate(flatten(final_list)) flat = to_list(res) assert len(flat) == 6 for i in range(6): assert count(flat[i]) == i + 1
def test_global_function(): mod = relay.Module() ident = relay.GlobalVar('ident') a = relay.TypeVar('a') v = relay.Var('v', a) mod[ident] = relay.Function([v], v, a, [a]) call1 = ident(relay.const(1)) call2 = ident(relay.Tuple([relay.const(2), relay.const(2)])) call_val1 = run_as_python(call1, mod) assert_tensor_value(call_val1, 1) call_val2 = run_as_python(call2, mod) assert_adt_len(call_val2, 2) assert_tensor_value(call_val2[0], 2) assert_tensor_value(call_val2[1], 2)
def test_concat(): a = relay.TypeVar("a") assert prelude.mod[concat].checked_type == relay.FuncType( [rlist(a), rlist(a)], rlist(a), [a]) l1 = cons(make_nat_expr(prelude, 1), cons(make_nat_expr(prelude, 2), nil())) l2 = cons(make_nat_expr(prelude, 3), cons(make_nat_expr(prelude, 4), nil())) res = eval(concat(l1, l2)) catted = to_list(res) assert len(catted) == 4 assert count(catted[0]) == 1 assert count(catted[1]) == 2 assert count(catted[2]) == 3 assert count(catted[3]) == 4
def test_filter(): a = relay.TypeVar("a") expected_type = relay.FuncType( [relay.FuncType([a], relay.scalar_type("bool")), rlist(a)], rlist(a), [a] ) assert prelude.mod[filter].checked_type == expected_type x = relay.Var("x", nat()) greater_than_one = relay.Function( [x], relay.Match( x, [ relay.Clause( relay.PatternConstructor( s, [relay.PatternConstructor(s, [relay.PatternWildcard()])] ), relay.const(True), ), relay.Clause(relay.PatternWildcard(), relay.const(False)), ], ), ) res = intrp.evaluate( filter( greater_than_one, cons( make_nat_expr(prelude, 1), cons( make_nat_expr(prelude, 1), cons( make_nat_expr(prelude, 3), cons( make_nat_expr(prelude, 1), cons(make_nat_expr(prelude, 5), cons(make_nat_expr(prelude, 1), nil())), ), ), ), ), ) ) filtered = to_list(res) assert len(filtered) == 2 assert count(filtered[0]) == 3 assert count(filtered[1]) == 5
def test_rev(): a = relay.TypeVar("a") assert prelude.mod[rev].checked_type == relay.FuncType([rlist(a)], rlist(a), [a]) res = intrp.evaluate( rev( cons( make_nat_expr(prelude, 1), cons(make_nat_expr(prelude, 2), cons(make_nat_expr(prelude, 3), nil())), ) ) ) reversed = to_list(res) assert len(reversed) == 3 assert count(reversed[0]) == 3 assert count(reversed[1]) == 2 assert count(reversed[2]) == 1
def test_global_recursion(): mod = tvm.IRModule() p = Prelude(mod) rlist, cons, nil = p.mod.get_type("List") copy = relay.GlobalVar("copy") # same as above: it copies the given list a = relay.TypeVar("a") v = relay.Var("v", rlist(a)) h = relay.Var("h") t = relay.Var("t") copy_def = relay.Function( [v], relay.Match( v, [ relay.Clause( relay.PatternConstructor( cons, [relay.PatternVar(h), relay.PatternVar(t)]), cons(h, copy(t)), ), relay.Clause(relay.PatternConstructor(nil, []), nil()), ], ), rlist(a), [a], ) mod[copy] = copy_def call1 = copy_def(cons(relay.const(1), cons(relay.const(2), nil()))) val1 = run_as_python(call1, mod) assert_constructor_value(val1, cons, 2) assert_tensor_value(val1.fields[0], 1) assert_constructor_value(val1.fields[1], cons, 2) assert_tensor_value(val1.fields[1].fields[0], 2) assert_constructor_value(val1.fields[1].fields[1], nil, 0) call2 = copy_def(cons(relay.Tuple([]), nil())) val2 = run_as_python(call2, mod) assert_constructor_value(val2, cons, 2) assert_adt_len(val2.fields[0], 0) assert_constructor_value(val2.fields[1], nil, 0)
def test_foldr1(): a = relay.TypeVar("a") lhs = prelude.mod[foldr1].checked_type rhs = relay.FuncType([relay.FuncType([a, a], a), rlist(a)], a, [a]) assert lhs == rhs x = relay.Var("x") y = relay.Var("y") f = relay.Function([x, y], add(x, y)) res = eval( foldr1( f, cons( make_nat_expr(prelude, 1), cons(make_nat_expr(prelude, 2), cons(make_nat_expr(prelude, 3), nil())), ), )) assert count(res) == 6
def test_head_cons(): mod = relay.Module() p = Prelude(mod) def hd_impl(): a = relay.TypeVar("a") x = relay.Var("x", p.l(a)) y = relay.Var("y") z = relay.Var("z") cons_case = relay.Clause(relay.PatternConstructor(p.cons, [relay.PatternVar(y), relay.PatternVar(z)]), y) return relay.Function([x], relay.Match(x, [cons_case]), a, [a]) t = relay.TypeVar("t") x = relay.Var("x", t) hd = relay.Var("hd") body = relay.Let(hd, hd_impl(), hd(p.cons(x, p.nil()))) f = relay.Function([x], body, None, [t]) f = infer_type(f, mod=mod) res = dcpe(f) assert alpha_equal(res, relay.Function([x], x, t, [t]))
def test_nested_matches(): a = relay.TypeVar('a') x = relay.Var('x') y = relay.Var('y') w = relay.Var('w') h = relay.Var('h') t = relay.Var('t') flatten = relay.GlobalVar('flatten') # flatten could be written using a fold, but this way has nested matches inner_match = relay.Match(y, [ relay.Clause(relay.PatternConstructor(nil), flatten(w)), relay.Clause( relay.PatternConstructor( cons, [relay.PatternVar(h), relay.PatternVar(t)]), cons(h, flatten(cons(t, w)))) ]) mod[flatten] = relay.Function( [x], relay.Match(x, [ relay.Clause(relay.PatternConstructor(nil), nil()), relay.Clause( relay.PatternConstructor( cons, [relay.PatternVar(y), relay.PatternVar(w)]), inner_match) ]), l(a), [a]) first_list = cons(build_nat(1), cons(build_nat(2), cons(build_nat(3), nil()))) second_list = cons(build_nat(4), cons(build_nat(5), cons(build_nat(6), nil()))) final_list = cons(first_list, cons(second_list, nil())) res = intrp.evaluate(flatten(final_list)) flat = to_list(res) assert len(flat) == 6 for i in range(6): assert count(flat[i]) == i + 1
def test_single_constructor_adt(): mod = tvm.IRModule() box = relay.GlobalTypeVar("box") a = relay.TypeVar("a") box_ctor = relay.Constructor("box", [a], box) box_data = relay.TypeData(box, [a], [box_ctor]) mod[box] = box_data v = relay.Var("v") match = relay.Match(v, [ relay.Clause( relay.PatternConstructor(box_ctor, [relay.PatternWildcard()]), v) ]) # with one constructor, having one pattern constructor case is exhaustive assert len(unmatched_cases(match, mod)) == 0 # this will be so if we nest the constructors too nested_pattern = relay.Match( v, [ relay.Clause( relay.PatternConstructor( box_ctor, [ relay.PatternConstructor( box_ctor, [ relay.PatternConstructor( box_ctor, [relay.PatternWildcard()]) ], ) ], ), v, ) ], ) assert len(unmatched_cases(nested_pattern, mod)) == 0
def init_box_adt(mod): box = relay.GlobalTypeVar('box') a = relay.TypeVar('a') box_ctor = relay.Constructor('box', [a], box) mod[box] = relay.TypeData(box, [a], [box_ctor]) return (box, box_ctor)