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)
