def test_recursion():
"""
Program:
def @f(%n: int32, %data: float32) -> float32 {
if (%n == 0) {
%data
} else {
@f(%n - 1, log(%data))
}
}
"""
sb = relay.ScopeBuilder()
f = relay.GlobalVar("f")
ti32 = relay.scalar_type("int32")
tf32 = relay.scalar_type("float32")
n = relay.var("n", ti32)
data = relay.var("data", tf32)
with sb.if_scope(relay.equal(n, relay.const(0, ti32))):
sb.ret(data)
with sb.else_scope():
sb.ret(f(relay.subtract(n, relay.const(1, ti32)), relay.log(data)))
mod = relay.Module()
mod[f] = relay.Function([n, data], sb.get())
assert "@f(%1, %2) /* ty=float32 */" in mod.astext()
assert mod[f].checked_type == relay.FuncType([ti32, tf32], tf32)
def test_equal():
i = relay.var('i', shape=[], dtype='int32')
eq = op.equal(i, relay.const(0, dtype='int32'))
func = relay.Function([i], eq)
ft = relay.ir_pass.infer_type(func)
assert ft.checked_type == relay.FuncType([relay.scalar_type('int32')], relay.scalar_type('bool'))
def test_monomorphic_let():
"Program: let %x = 1; %x"
sb = relay.ScopeBuilder()
x = sb.let('x', relay.const(1.0, "float64"))
sb.ret(x)
xchecked = relay.ir_pass.infer_type(sb.get())
assert xchecked.checked_type == relay.scalar_type("float64" )
def test_filter():
a = relay.TypeVar("a")
expected_type = relay.FuncType([
relay.FuncType([a], relay.scalar_type("bool")), l(a)
], l(a), [a])
assert 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(build_nat(1),
cons(build_nat(1),
cons(build_nat(3),
cons(build_nat(1),
cons(build_nat(5),
cons(build_nat(1),
nil()))))))))
filtered = to_list(res)
assert len(filtered) == 2
assert count(filtered[0]) == 3
assert count(filtered[1]) == 5
def test_incomplete_call():
tt = relay.scalar_type('int32')
x = relay.var('x', tt)
f = relay.var('f')
func = relay.Function([x, f], relay.Call(f, [x]), tt)
ft = relay.ir_pass.infer_type(func)
f_type = relay.FuncType([tt], tt)
assert ft.checked_type == relay.FuncType([tt, f_type], tt)
def verify_full(fill_value, src_shape, dtype):
x = relay.var("x", relay.scalar_type(dtype))
z = relay.full(x, src_shape, dtype)
func = relay.Function([x], z)
ref_res = np.full(src_shape, fill_value)
for target, ctx in ctx_list():
for kind in ["graph", "debug"]:
intrp = relay.create_executor(kind, ctx=ctx, target=target)
op_res = intrp.evaluate(func)(np.array(fill_value, dtype))
tvm.testing.assert_allclose(op_res.asnumpy(), ref_res, rtol=1e-5)
def test_ref():
x = relay.var("x", "float32")
y = relay.var("y", "float32")
r = relay.RefCreate(x)
st = relay.scalar_type("float32")
assert relay.ir_pass.infer_type(r).checked_type == relay.RefType(st)
g = relay.RefRead(r)
assert relay.ir_pass.infer_type(g).checked_type == st
w = relay.RefWrite(r, y)
assert relay.ir_pass.infer_type(w).checked_type == relay.TupleType([])
def test_global_var_recursion():
mod = relay.Module({})
gv = relay.GlobalVar("foo")
x = relay.var('x', shape=[])
tt = relay.scalar_type('float32')
func = relay.Function([x], relay.Call(gv, [x]), tt)
mod[gv] = func
ft = relay.ir_pass.infer_type(gv, mod)
assert mod[ft].checked_type == relay.FuncType([tt], tt)
def test_id_type():
mod = tvm.IRModule()
id_type = relay.GlobalTypeVar("id")
a = relay.TypeVar("a")
mod[id_type] = relay.TypeData(id_type, [a], [])
b = relay.TypeVar("b")
make_id = relay.Var("make_id", relay.FuncType([b], id_type(b), [b]))
t = relay.scalar_type("float32")
b = relay.Var("b", t)
mod["main"] = relay.Function([make_id, b], make_id(b))
mod = transform.InferType()(mod)
assert mod["main"].body.checked_type == id_type(t)
def verify_full(fill_value, src_shape, dtype):
x = relay.var("x", relay.scalar_type(dtype))
rank = len(src_shape)
dyn_src_shape = relay.var("dyn_scr_shape",
relay.ty.TensorType((rank, ), "int64"))
z = relay.full(x, dyn_src_shape, dtype)
func = relay.Function([x, dyn_src_shape], z)
ref_res = np.full(src_shape, fill_value).astype(dtype)
verify_func(func, [
np.array(fill_value).astype(dtype),
np.array(src_shape).astype("int64")
], ref_res)
def verify_full_like(base, fill_value, dtype):
x_data = np.random.uniform(low=-1, high=1, size=base).astype(dtype)
x = relay.var("x", relay.TensorType(base, dtype))
y = relay.var("y", relay.scalar_type(dtype))
z = relay.full_like(x, y)
func = relay.Function([x, y], z)
ref_res = np.full_like(x_data, fill_value)
for target, ctx in ctx_list():
for kind in ["graph", "debug"]:
intrp = relay.create_executor(kind, ctx=ctx, target=target)
op_res = intrp.evaluate(func)(x_data, np.array(fill_value, dtype))
tvm.testing.assert_allclose(op_res.asnumpy(), ref_res, rtol=1e-5)
def verify_full_like(base, fill_value, dtype):
x_data = np.random.uniform(low=-1, high=1, size=base).astype(dtype)
x = relay.var("x", relay.TensorType(base, dtype))
y = relay.var("y", relay.scalar_type(dtype))
z = relay.full_like(x, y)
func = relay.Function([x, y], z)
ref_res = np.full_like(x_data, fill_value)
for target, ctx in ctx_list():
for kind in ["graph", "debug"]:
intrp = relay.create_executor(kind, ctx=ctx, target=target)
op_res = intrp.evaluate(func)(x_data, np.array(fill_value, dtype))
tvm.testing.assert_allclose(op_res.asnumpy(), ref_res, rtol=1e-5)
def verify_full(fill_value, fill_shape, dtype):
x = relay.var("x", relay.scalar_type(dtype))
y = relay.var("y", relay.TensorType(fill_shape, 'int64'))
z = relay.full(x, relay.shape_of(y), dtype)
func = run_infer_type(relay.Function([x, y], z))
func2 = run_opt_pass(run_opt_pass(func, transform.DynamicToStatic()), transform.InferType())
zz = func2.body
assert isinstance(zz, relay.Call)
assert zz.op == relay.op.get("full")
ref_res = np.full(fill_shape, fill_value).astype(dtype)
y_data = np.random.uniform(low=-1, high=1, size=fill_shape).astype('int64')
verify_func(func2, [fill_value, y_data], ref_res)
def test_higher_order_nested():
a = relay.TypeVar("a")
x = relay.Var("x", a)
id_func = relay.Function([x], x, a, [a])
choice_t = relay.FuncType([], relay.scalar_type("bool"))
f = relay.Var("f", choice_t)
b = relay.TypeVar("b")
z = relay.Var("z")
top = relay.Function([f], relay.If(f(), id_func, relay.Function([z], z)),
relay.FuncType([b], b), [b])
expected = relay.FuncType([choice_t], relay.FuncType([b], b), [b])
ft = infer_expr(top)
assert ft.checked_type == expected
def test_higher_order_nested():
a = relay.TypeVar('a')
x = relay.Var('x', a)
id_func = relay.Function([x], x, a, [a])
choice_t = relay.FuncType([], relay.scalar_type('bool'))
f = relay.Var('f', choice_t)
b = relay.TypeVar('b')
z = relay.Var('z')
top = relay.Function([f], relay.If(f(), id_func, relay.Function([z], z)),
relay.FuncType([b], b), [b])
expected = relay.FuncType([choice_t], relay.FuncType([b], b), [b])
ft = run_infer_type(top)
assert ft.checked_type == expected
def test_higher_order_argument():
a = relay.TypeVar('a')
x = relay.Var('x', a)
id_func = relay.Function([x], x, a, [a])
b = relay.TypeVar('b')
f = relay.Var('f', relay.FuncType([b], b))
y = relay.Var('y', b)
ho_func = relay.Function([f, y], f(y), b, [b])
# id func should be an acceptable argument to the higher-order
# function even though id_func takes a type parameter
ho_call = ho_func(id_func, relay.const(0, 'int32'))
hc = relay.ir_pass.infer_type(ho_call)
expected = relay.scalar_type('int32')
assert hc.checked_type == expected
def test_higher_order_argument():
a = relay.TypeVar("a")
x = relay.Var("x", a)
id_func = relay.Function([x], x, a, [a])
b = relay.TypeVar("b")
f = relay.Var("f", relay.FuncType([b], b))
y = relay.Var("y", b)
ho_func = relay.Function([f, y], f(y), b, [b])
# id func should be an acceptable argument to the higher-order
# function even though id_func takes a type parameter
ho_call = ho_func(id_func, relay.const(0, "int32"))
hc = infer_expr(ho_call)
expected = relay.scalar_type("int32")
assert hc.checked_type == expected
def test_higher_order_argument():
a = relay.TypeVar('a')
x = relay.Var('x', a)
id_func = relay.Function([x], x, a, [a])
b = relay.TypeVar('b')
f = relay.Var('f', relay.FuncType([b], b))
y = relay.Var('y', b)
ho_func = relay.Function([f, y], f(y), b, [b])
# id func should be an acceptable argument to the higher-order
# function even though id_func takes a type parameter
ho_call = ho_func(id_func, relay.const(0, 'int32'))
hc = run_infer_type(ho_call)
expected = relay.scalar_type('int32')
assert hc.checked_type == expected
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_higher_order_nested():
a = relay.TypeVar('a')
x = relay.Var('x', a)
id_func = relay.Function([x], x, a, [a])
choice_t = relay.FuncType([], relay.scalar_type('bool'))
f = relay.Var('f', choice_t)
b = relay.TypeVar('b')
z = relay.Var('z')
top = relay.Function(
[f],
relay.If(f(), id_func, relay.Function([z], z)),
relay.FuncType([b], b),
[b])
expected = relay.FuncType([choice_t], relay.FuncType([b], b), [b])
ft = relay.ir_pass.infer_type(top)
assert ft.checked_type == expected
def verify_meshgrid(lengths, indexing='ij'):
input_vars = []
input_data = []
for (i, length) in enumerate(lengths):
input_name = 'x_{}'.format(i)
if (length == 0):
input_vars.append(relay.var(input_name, relay.scalar_type('float32')))
input_data.append(np.array(1, 'float32'))
else:
input_vars.append(relay.var(input_name, relay.TensorType((length,), 'float32')))
input_data.append(np.arange(length).astype('float32'))
z = relay.meshgrid(input_vars, indexing=indexing).astuple()
func = relay.Function(input_vars, z)
ref_res = np.meshgrid(*input_data, indexing=indexing)
for (target, ctx) in tvm.testing.enabled_targets():
for kind in ['graph', 'debug']:
intrp = relay.create_executor(kind, ctx=ctx, target=target)
op_res = intrp.evaluate(func)(*input_data)
assert (len(op_res) == len(ref_res))
for i in range(len(op_res)):
tvm.testing.assert_allclose(op_res[i].asnumpy(), ref_res[i], rtol=1e-05)
def test_sum():
assert mod[sum].checked_type == relay.FuncType(
[l(relay.scalar_type('int32'))], relay.scalar_type('int32'))
res = intrp.evaluate(sum(cons(relay.const(1), cons(relay.const(2),
nil()))))
assert get_scalar(res) == 3
def test_length():
a = relay.TypeVar("a")
assert mod[length].checked_type == relay.FuncType(
[l(a)], relay.scalar_type('int32'), [a])
res = intrp.evaluate(length(cons(z(), cons(z(), cons(z(), nil())))))
assert get_scalar(res) == 3
def test_sum():
assert prelude.mod[sum].checked_type == relay.FuncType(
[rlist(relay.scalar_type("int32"))], relay.scalar_type("int32")
)
res = intrp.evaluate(sum(cons(relay.const(1), cons(relay.const(2), nil()))))
assert get_scalar(res) == 3
from tvm import relay
import tvm
x = relay.expr.var('x', relay.scalar_type('float32'), dtype='float32')
y = relay.expr.var('y', relay.scalar_type('float32'), dtype='float32')
dense = relay.nn.dense(x, y)
func = relay.expr.Function([x], dense, relay.scalar_type('float32'))
mod = relay.Module.from_expr(func) # note this API
print("Relay module function:\n", mod.astext(show_meta_data=False))
graph, lib, params = tvm.relay.build(mod, 'llvm', params={})
print("TVM graph:\n", graph)
print("TVM parameters:\n", params)
print("TVM compiled target function:\n", lib.get_source())
"bool",
"int8x4",
"uint1x4",
"float16x4",
}
def parses_as(code, expr):
# type: (str, relay.Expr) -> bool
return alpha_equal(relay.fromtext(SEMVER + "\n" + code), expr)
def get_scalar(x):
# type: (relay.Constant) -> (Union[float, int, bool])
return x.data.asnumpy().item()
int32 = relay.scalar_type("int32")
_ = relay.Var("_")
X = relay.Var("x")
Y = relay.Var("y")
X_ANNO = relay.Var("x", int32)
Y_ANNO = relay.Var("y", int32)
UNIT = relay.Tuple([])
# decorator to determine if parser is enabled
def if_parser_enabled(func):
# https://stackoverflow.com/q/7727678
@wraps(func)
def wrapper():
if not enabled():
from tvm import relay
import tvm.relay.op
x = relay.expr.var('x', relay.scalar_type('int64'), dtype= 'int64')
one = relay.expr.const(1, dtype= 'int64')
add = relay.op.tensor.add(x, one)
func = relay.Function([x], add, relay.scalar_type('int64'))
mod = tvm.ir.IRModule.from_expr(func)
graph, lib, params = tvm.relay.build(mod, 'llvm', params={})
print("TVM graph:\n", graph)
print("TVM parameters:\n", params)
print("TVM compiled target function:\n", lib.get_source())
return expr
def parses_as(code, expr):
# type: (str, relay.Expr) -> bool
parsed = parse_text(code)
result = graph_equal(parsed, expr)
return result
def get_scalar(x):
# type: (relay.Constant) -> (Union[float, int, bool])
return x.data.asnumpy().item()
int32 = relay.scalar_type("int32")
_ = relay.Var("_")
X = relay.Var("x")
Y = relay.Var("y")
X_ANNO = relay.Var("x", int32)
Y_ANNO = relay.Var("y", int32)
UNIT = relay.Tuple([])
def test_comments():
assert parses_as(
"""
// This is a line comment!
()
def test_free_expr():
x = relay.var("x", "float32")
y = relay.add(x, x)
yy = relay.ir_pass.infer_type(y)
assert yy.checked_type == relay.scalar_type("float32")
assert x.vid.same_as(yy.args[0].vid)
def test_free_expr():
x = relay.var("x", "float32")
y = relay.add(x, x)
yy = relay.ir_pass.infer_type(y)
assert yy.checked_type == relay.scalar_type("float32")
assert x.vid.same_as(yy.args[0].vid)
def test_length():
a = relay.TypeVar("a")
assert prelude.mod[length].checked_type == relay.FuncType(
[rlist(a)], relay.scalar_type("int32"), [a])
res = eval(length(cons(z(), cons(z(), cons(z(), nil())))))
assert get_scalar(res) == 3