def test_pow():
mod = tvm.IRModule()
p = Prelude(mod)
p.mod.import_from_std("nat.rly")
nat_iterate = mod.get_global_var("nat_iterate")
shape = (10, 10)
dtype = "float32"
t = relay.TensorType(shape, dtype)
x = relay.var("x", t)
double = relay.Function([x], x + x)
i = relay.var("i", t)
func = relay.Function([i], nat_iterate(double, make_nat_expr(p, 3))(i))
mod["main"] = func
mod = transform.InferType()(mod)
mod["main"] = gradient(mod["main"], mod=mod)
m = transform.InferType()(mod)
back_func = m["main"]
assert back_func.checked_type == relay.FuncType(
[t], relay.TupleType([t, relay.TupleType([t])]))
i_nd = rand(dtype, *shape)
forward, (grad_i, ) = create_executor(mod=mod).evaluate(back_func)(i_nd)
tvm.testing.assert_allclose(forward.numpy(), 8 * i_nd.numpy())
tvm.testing.assert_allclose(grad_i.numpy(),
8 * np.ones_like(grad_i.numpy()))
def test_tuple():
ttype = relay.TupleType([relay.TensorType((1,)), relay.TensorType((10,))])
tup = relay.var('tup', type_annotation=ttype)
f = relay.Function([tup], relay.TupleGetItem(tup, 1))
i_data = np.random.rand(41).astype('float32')
j_data = np.random.rand(10).astype('float32')
exe = create_exec(f)
code, lib = exe.save()
des_exec = _vm.Executable.load_exec(code, lib)
des_vm = _vm.VirtualMachine(des_exec)
des_vm.init(tvm.cpu())
result = veval(des_vm, (i_data, j_data))
tvm.testing.assert_allclose(result.asnumpy(), j_data)
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_function():
param_names = ['a', 'b', 'c', 'd']
params = tvm.runtime.convert([relay.Var(n) for n in param_names])
ret_type = relay.TupleType(tvm.runtime.convert([]))
body = relay.Tuple(tvm.runtime.convert([]))
type_params = tvm.runtime.convert([])
fn = relay.Function(params, body, ret_type, type_params)
fn = fn.with_attr("test_attribute", tvm.tir.StringImm("value"))
assert fn.params == params
assert fn.body == body
assert fn.type_params == type_params
assert fn.span == None
assert fn.attrs["test_attribute"] == "value"
str(fn)
check_json_roundtrip(fn)
def test_vm_reshape_tuple(target, dev, x_shape=(1, 4, 2), y_shape=(1, 2, 10)):
tup = relay.var(
"tup",
type_annotation=relay.TupleType(
[relay.TensorType(x_shape),
relay.TensorType(y_shape)]),
)
out = relay.reshape(relay.TupleGetItem(tup, 0), (1, -1))
f = relay.Function([tup], out)
x_data = np.random.uniform(size=x_shape).astype("float32")
y_data = np.random.uniform(size=y_shape).astype("float32")
res = veval(f, (x_data, y_data), device=dev, target=target)
tvm.testing.assert_allclose(res.numpy(), np.reshape(x_data, (1, -1)))
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_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_tuple():
ttype = relay.TupleType([relay.TensorType((1,)), relay.TensorType((10,))])
tup = relay.var('tup', type_annotation=ttype)
f = relay.Function([tup], relay.TupleGetItem(tup, 1))
i_data = np.random.rand(41).astype('float32')
j_data = np.random.rand(10).astype('float32')
vm = create_vm(f)
ser = serializer.Serializer(vm)
code, lib = ser.serialize()
deser = deserializer.Deserializer(code, lib)
des_vm = deser.deserialize()
result = veval(des_vm, (i_data, j_data))
tvm.testing.assert_allclose(result.asnumpy(), j_data)
def test_func_kind():
# only contain type kinds
tp1 = relay.TypeParam('tp1', relay.Kind.Type)
tp2 = relay.TypeParam('tp2', relay.Kind.Type)
shape = tvm.convert([1, 2, 3])
dtype = 'float32'
tensor_type = relay.TensorType(shape, dtype)
type_params = tvm.convert([tp1, tp2])
type_constraints = tvm.convert([])
arg_types = tvm.convert([tp1, tensor_type])
ret_type = relay.TupleType(tvm.convert([tp2, tensor_type]))
tf = relay.FuncType(arg_types, ret_type, type_params, type_constraints)
assert check_kind(tf)
def test_adt_match():
mod = relay.Module()
box, constructor = initialize_box_adt(mod)
v = relay.Var('v', relay.TensorType((), 'float32'))
match = relay.Match(constructor(relay.const(0, 'float32')),
[relay.Clause(
relay.PatternConstructor(constructor,
[relay.PatternVar(v)]),
relay.Tuple([])),
# redundant but shouldn't matter to typechecking
relay.Clause(relay.PatternWildcard(),
relay.Tuple([]))])
mt = relay.ir_pass.infer_type(match, mod)
assert mt.checked_type == relay.TupleType([])
def test_function():
param_names = ["a", "b", "c", "d"]
params = tvm.runtime.convert([relay.Var(n) for n in param_names])
ret_type = relay.TupleType(tvm.runtime.convert([]))
body = relay.Tuple(tvm.runtime.convert([]))
type_params = tvm.runtime.convert([])
fn = relay.Function(params, body, ret_type, type_params)
fn = fn.with_attr("test_attribute", "value")
fn = fn.with_attr("test_attribute1", "value1")
assert fn.params == params
assert fn.body == body
assert fn.type_params == type_params
assert fn.span == None
assert fn.attrs["test_attribute"] == "value"
assert fn.attrs["test_attribute1"] == "value1"
str(fn)
check_json_roundtrip(fn)
def test_func_kind():
# only contain type kinds
tp1 = relay.TypeVar('tp1', relay.TypeKind.Type)
tp2 = relay.TypeVar('tp2', relay.TypeKind.Type)
shape = tvm.runtime.convert([1, 2, 3])
dtype = 'float32'
tensor_type = relay.TensorType(shape, dtype)
tr = relay.TypeRelation(None, tvm.runtime.convert([tensor_type, tp1]) , 1, None)
type_params = tvm.runtime.convert([tp1, tp2])
type_constraints = tvm.runtime.convert([tr])
arg_types = tvm.runtime.convert([tp1, tensor_type])
ret_type = relay.TupleType(tvm.runtime.convert([tp2, tensor_type]))
tf = relay.FuncType(arg_types, ret_type, type_params, type_constraints)
assert check_kind(tf) == relay.TypeKind.Type
def test_adt_match_type_annotations():
mod = relay.Module()
box, constructor = initialize_box_adt(mod)
# the only type annotation is inside the match pattern var
# but that should be enough info
tt = relay.TensorType((2, 2), 'float32')
x = relay.Var('x')
mv = relay.Var('mv', tt)
match = relay.Match(constructor(x), [
relay.Clause(
relay.PatternConstructor(constructor, [relay.PatternVar(mv)]),
relay.Tuple([]))
])
func = relay.Function([x], match)
ft = relay.ir_pass.infer_type(func, mod)
assert ft.checked_type == relay.FuncType([tt], relay.TupleType([]))
def test_constructor_call():
mod = tvm.IRModule()
box, constructor = initialize_box_adt(mod)
box_unit = constructor(relay.Tuple([]))
box_constant = constructor(relay.const(0, "float32"))
func = relay.Function([], relay.Tuple([box_unit, box_constant]))
mod["main"] = func
mod = infer_mod(mod)
ret_type = mod["main"].checked_type.ret_type.fields
# NB(@jroesch): when we annotate spans the ast fragments before
# annotation the previous fragments will no longer be directly equal.
box = mod.get_global_type_var("box")
expected1 = box(relay.TupleType([]))
expected2 = box(relay.TensorType((), "float32"))
assert ret_type[0] == expected1
assert ret_type[1] == expected2
def test_zip():
a = relay.TypeVar("a")
b = relay.TypeVar("b")
expected_type = relay.FuncType([rlist(a), rlist(b)],
rlist(relay.TupleType([a, b])), [a, b])
assert prelude.mod[zip].checked_type == expected_type
l1 = cons(
make_nat_expr(prelude, 1),
cons(make_nat_expr(prelude, 2), cons(make_nat_expr(prelude, 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(prelude, 4), cons(make_nat_expr(prelude, 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 override_shape(tensor_type, axis, dim):
"""Change a dimension in a tensor shape."""
# Handle multiple tensors by overriding the shape of each.
if isinstance(tensor_type, relay.TupleType):
tensor_type = tensor_type.fields
else:
tensor_type = [tensor_type]
# Create new tensortypes for each input.
new_types = []
for t_type in tensor_type:
new_dims = list(t_type.shape)
new_dims[axis] = dim
new_types.append(relay.TensorType(new_dims, t_type.dtype))
# Dont return a tuple if there is a single tensor.
if len(new_types) == 1:
return new_types[0]
return relay.TupleType(tvm.runtime.convert(new_types))
def _test_tuple_argument(mode):
shape = (2, 3)
dtype = "float32"
tensor_type = relay.TensorType(shape, dtype)
fields = 3
tuple_type = relay.TupleType([tensor_type] * fields)
tup = relay.var("tup", type_annotation=tuple_type)
body = relay.TupleGetItem(tup, 0)
for i in range(1, fields):
body = relay.add(body, relay.TupleGetItem(tup, i))
func = relay.Function([tup], body)
func = run_infer_type(func)
back_func = run_infer_type(gradient(func, mode=mode))
xs = [rand(dtype, *shape) for _ in range(fields)]
xs_np = np.array([x.numpy() for x in xs])
expected_forward = np.sum(xs_np, axis=0)
forward, grad = create_executor().evaluate(back_func)(tuple(xs))
tvm.testing.assert_allclose(forward.numpy(), expected_forward)
for field in grad[0]:
tvm.testing.assert_allclose(field.numpy(), np.ones_like(field.numpy()))
def test_adt_match():
mod = tvm.IRModule()
box, constructor = initialize_box_adt(mod)
v = relay.Var("v", relay.TensorType((), "float32"))
match = relay.Match(
constructor(relay.const(0, "float32")),
[
relay.Clause(
relay.PatternConstructor(constructor, [relay.PatternVar(v)]),
relay.Tuple([])),
# redundant but shouldn't matter to typechecking
relay.Clause(relay.PatternWildcard(), relay.Tuple([])),
],
)
func = relay.Function([], match)
mod["main"] = func
mod = infer_mod(mod)
actual = mod["main"].checked_type.ret_type
assert actual == relay.TupleType([])
def test_adt_match_type_annotations():
mod = tvm.IRModule()
box, constructor = initialize_box_adt(mod)
# the only type annotation is inside the match pattern var
# but that should be enough info
tt = relay.TensorType((2, 2), "float32")
x = relay.Var("x")
mv = relay.Var("mv", tt)
match = relay.Match(
constructor(x),
[
relay.Clause(
relay.PatternConstructor(constructor, [relay.PatternVar(mv)]),
relay.Tuple([]))
],
)
mod["main"] = relay.Function([x], match)
mod = infer_mod(mod)
ft = mod["main"].checked_type
assert ft == relay.FuncType([tt], relay.TupleType([]))
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_add_tuple():
"""Add elements of tuple. Check types and semantic equivalence."""
mod = tvm.IRModule()
shape = (10, 10)
dtype = "float32"
tensor_type = relay.TensorType(shape, dtype)
t = relay.TupleType([tensor_type, tensor_type])
x = relay.var("x", t)
# f((x1,x2)) = x1 + x2
y = relay.Function([x], relay.TupleGetItem(x, 0) + relay.TupleGetItem(x, 1))
mod["main"] = y
mod = transform.InferType()(mod)
mod = transform.LazyGradientInit()(mod)
mod = tvm.transform.PrintIR(show_meta_data=True)(mod)
y = mod["main"]
assert mod["main"].checked_type == relay.FuncType([t], tensor_type)
x = (rand(dtype, *shape), rand(dtype, *shape))
y = create_executor(mod=mod).evaluate(y)(x)
assert_allclose(y.numpy(), x[0].numpy() + x[1].numpy())
def test_ret_tuple():
"""Test tuple return type. Check types and semantic equivalence."""
mod = tvm.IRModule()
shape = (10, 10)
dtype = "float32"
t = relay.TensorType(shape, dtype)
x = relay.var("x", t)
# f(x) = (x,x)
func = relay.Function([x], relay.Tuple([x, x * relay.const(2.0)]))
func = run_infer_type(func)
mod["main"] = func
mod = transform.InferType()(mod)
mod = transform.LazyGradientInit()(mod)
func = mod["main"]
assert mod["main"].checked_type == relay.FuncType([t], relay.TupleType([t, t]))
x = rand(dtype, *shape)
y = create_executor(mod=mod).evaluate(func)(x)
assert_allclose(y[0].numpy(), x.numpy())
assert_allclose(y[1].numpy(), x.numpy() * 2.0)
def test_function_alpha_equal():
tt1 = relay.TensorType((1, 2, 3), "float32")
tt2 = relay.TensorType((4, 5, 6), "int8")
tt3 = relay.TupleType([tt1, tt2])
v1 = relay.Var("v1", tt1)
v2 = relay.Var("v2", tt2)
v3 = relay.Var("v3", tt3)
v4 = relay.Var("v4", tt2)
vret = relay.Constant(tvm.nd.array(np.ones(1)))
tp1 = relay.TypeVar("tp1", relay.Kind.Type)
tp2 = relay.TypeVar("tp2", relay.Kind.Type)
tp3 = relay.TypeVar("tp3", relay.Kind.Shape)
tp4 = relay.TypeVar("tp4", relay.Kind.Shape)
basic_args = [relay.Var("v3", tt1), relay.Var("v4", tt2)]
basic_tps = [tp1, tp2]
func = relay.Function([v1, v2], v1,
tt2, basic_tps)
mapped = relay.Function(basic_args, basic_args[0], tt2, basic_tps)
assert alpha_equal(func, mapped)
fewer_params = relay.Function([relay.Var("v4", tt2)], v4, tt2, basic_tps)
assert not alpha_equal(func, fewer_params)
more_params = relay.Function([relay.Var("v3", tt1),
relay.Var("v4", tt2),
relay.Var("v2", tt2)], v4, tt2, basic_tps)
assert not alpha_equal(func, more_params)
params_unordered = relay.Function([v2, v1], v1,
tt2, basic_tps)
assert not alpha_equal(func, params_unordered)
params_mismatch = relay.Function([v1, v3], v1,
tt2, basic_tps)
assert not alpha_equal(func, params_mismatch)
# also would not typecheck
ret_type_mismatch = relay.Function(basic_args, v4, tt1, basic_tps)
assert not alpha_equal(func, ret_type_mismatch)
# also mis-typed
different_body = relay.Function(basic_args, v3, tt2, basic_tps)
assert not alpha_equal(func, different_body)
fewer_type_params = relay.Function(basic_args, v4, tt2, [tp1])
assert not alpha_equal(func, fewer_type_params)
more_type_params = relay.Function(basic_args, v4, tt2, [tp1, tp2, tp3])
assert not alpha_equal(func, more_type_params)
type_params_unordered = relay.Function(basic_args, v4, tt2, [tp2, tp1])
assert not alpha_equal(func, type_params_unordered)
different_type_params = relay.Function(basic_args, v4, tt2, [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, relay.Tuple([v3, v4]), tt3)
assert not alpha_equal(func, tupled_example)
# nullable
no_ret_type = relay.Function(basic_args, v4, None, [tp1, tp2])
# both null
assert alpha_equal(no_ret_type, no_ret_type)
# one null
assert not alpha_equal(func, no_ret_type)
assert not alpha_equal(no_ret_type, func)
def main(argv):
dtype = 'float32'
num_hidden = int(argv[1])
batch_size = 1
input_type = relay.TensorType((batch_size, num_hidden), dtype)
state_type = relay.TupleType([input_type, input_type])
weight_type = relay.TensorType((4 * num_hidden, num_hidden), dtype)
bias_type = relay.TensorType((4 * num_hidden, ), dtype)
# inputs = relay.Var('inputs', input_type)
# states = relay.Var('states', state_type)
# cell_state = relay.Var('cell_state', input_type)
# hidden_state = relay.Var('hidden_state', input_type)
# i2h_weight = relay.Var('i2h_weight', weight_type)
# i2h_bias = relay.Var('i2h_bias', bias_type)
# h2h_weight = relay.Var('h2h_weight', weight_type)
# h2h_bias = relay.Var('h2h_bias', bias_type)
# mod = tvm.IRModule()
# mod['lstm'] = lstm_cell(num_hidden)
# mod['main'] = relay.Function([inputs, cell_state, hidden_state,
# i2h_weight, i2h_bias, h2h_weight, h2h_bias],
# mod.get_global_var('lstm')(inputs, relay.Tuple([cell_state, hidden_state]),
# i2h_weight, i2h_bias, h2h_weight, h2h_bias))
mod, p = get_workload(batch_size, num_hidden)
ex = relay.create_executor('vm', mod=mod, ctx=tvm.cpu(), target='llvm')
i_val = generate_random_tensor(input_type)
cell_val = np.zeros((batch_size, num_hidden), np.float32)
hidden_val = np.zeros((batch_size, num_hidden), np.float32)
i2h_w_val = generate_random_tensor(weight_type)
i2h_b_val = generate_random_tensor(bias_type)
h2h_w_val = generate_random_tensor(weight_type)
h2h_b_val = generate_random_tensor(bias_type)
# order: i_sz, o_sz, input, cell, hidden, i2h_weight, h2h_weight, i2h_bias, h2h_bias
f = open(argv[2], 'wb')
f.write(num_hidden.to_bytes(4, 'little'))
f.write(num_hidden.to_bytes(4, 'little'))
i_val.asnumpy().tofile(f)
cell_val.tofile(f)
hidden_val.tofile(f)
i2h_w_val.asnumpy().tofile(f)
h2h_w_val.asnumpy().tofile(f)
i2h_b_val.asnumpy().tofile(f)
h2h_b_val.asnumpy().tofile(f)
print("Wrote %d bytes" % f.tell())
print("inputs:", i_val)
print("cell:", cell_val)
print("hidden:", hidden_val)
print("i2h_weights:", i2h_w_val)
print("h2h_weights:", h2h_w_val)
print("i2h_bias:", i2h_b_val)
print("h2h_bias:", h2h_b_val)
# i2h_dense = np.add(i2h_w_val.asnumpy().dot(i_val.asnumpy()[0]), i2h_b_val.asnumpy())
# h2h_dense = np.add(h2h_w_val.asnumpy().dot(hidden_val[0]), h2h_b_val.asnumpy())
# print("i2h dense: ", i2h_dense)
# print("h2h dense: ", h2h_dense)
# comb_dense = np.add(i2h_dense, h2h_dense)
# print("combined dense:", comb_dense)
# def sig(x):
# return (1.0 / (1.0 + math.exp(-x)))
# vsig = np.vectorize(sig)
# in_gate = vsig(comb_dense[:num_hidden])
# forget_gate = vsig(comb_dense[num_hidden:num_hidden*2])
# in_trx = np.tanh(comb_dense[num_hidden*2:num_hidden*3])
# out_gate = vsig(comb_dense[num_hidden*3:])
# next_c = np.add(np.multiply(forget_gate, cell_val), np.multiply(in_gate, in_trx))
# next_h = np.multiply(out_gate, np.tanh(next_c))
# print("next_c:", next_c)
# print("next_h:", next_h)
out = ex.evaluate()(i_val, i2h_w_val, i2h_b_val, h2h_w_val, h2h_b_val)
print("output: ", out)
out.asnumpy().tofile(f)
print("Wrote %d bytes" % f.tell())
f.close()
def test_unify_quantified_funcs_var_order():
solver = make_solver()
a, b, c = relay.TypeVar('a'), relay.TypeVar('b'), relay.TypeVar('c')
ft1 = relay.FuncType([a, relay.TupleType([b, c])], relay.TupleType([a, b, c]), [a, b, c])
ft2 = relay.FuncType([a, relay.TupleType([a, c])], relay.TupleType([a, a, c]), [a, c])
def lstm_definition(batch_size,
input_size,
hidden_size,
time_steps,
time_axis=1):
state_tensor_type = relay.TensorType((batch_size, hidden_size))
state_tuple_type = relay.TupleType([state_tensor_type, state_tensor_type])
input_var = relay.var("input", shape=(batch_size, time_steps, input_size))
state_var = relay.var("state", type_annotation=state_tuple_type)
i2h_weight_var = relay.var("i2h_weight",
shape=(4 * hidden_size, input_size))
h2h_weight_var = relay.var("h2h_weight",
shape=(4 * hidden_size, hidden_size))
i2h_bias_var = relay.var("i2h_bias", shape=(4 * hidden_size, ))
h2h_bias_var = relay.var("h2h_bias", shape=(4 * hidden_size, ))
# in this case, we are ignoring the state outputs
builder = relay.ScopeBuilder()
cell_var = builder.let(
"lstm_cell", relay_lstm_cell(batch_size, input_size, hidden_size))
splits = builder.let(
"splits",
relay.split(input_var, time_steps, time_axis).astuple())
last_state = state_var
seq_outs = []
for i in range(time_steps):
squeezed = builder.let(
f"squeezed_{i}",
relay.squeeze(relay.TupleGetItem(splits, i), axis=[time_axis]))
cell_out = builder.let(
f"cell_out_{i}",
cell_var(squeezed, last_state, i2h_weight_var, h2h_weight_var,
i2h_bias_var, i2h_bias_var))
new_seq_out = builder.let(f"seq_out_{i}",
relay.TupleGetItem(cell_out, 0))
seq_outs.append(new_seq_out)
new_hidden = builder.let(f"state_update_{i}",
relay.TupleGetItem(cell_out, 1))
last_state = new_hidden
stacked = builder.let("stacked", relay.stack(seq_outs, axis=time_axis))
# finally reshape to match pytorch's semantics (one layer)
reshape_hidden = builder.let(
"final_hidden",
relay.reshape(relay.TupleGetItem(last_state, 0),
(1, batch_size, hidden_size)))
reshape_cell = builder.let(
"final_cell",
relay.reshape(relay.TupleGetItem(last_state, 1),
(1, batch_size, hidden_size)))
builder.ret(relay.Tuple([stacked, reshape_hidden, reshape_cell]))
ret_type = relay.TupleType([
relay.TensorType((batch_size, time_steps, hidden_size)),
relay.TensorType((1, batch_size, hidden_size)),
relay.TensorType((1, batch_size, hidden_size))
])
return relay.Function([
input_var, state_var, i2h_weight_var, h2h_weight_var, i2h_bias_var,
h2h_bias_var
],
builder.get(),
ret_type=ret_type)
def test_tuple():
tp = relay.TensorType((10, ))
x = relay.var("x", tp)
res = relay.Tuple([x, x])
assert (relay.ir_pass.infer_type(res).checked_type == relay.TupleType(
[tp, tp]))
def lstm_cell(num_hidden, batch_size=1, dtype="float32", name=""):
"""Long-Short Term Memory (LSTM) network cell.
Parameters
----------
num_hidden : int
Number of units in output symbol.
batch_size : int
Batch size (length of states).
Returns
-------
result : tvm.relay.Function
A Relay function that evaluates an LSTM cell.
The function takes in a tensor of input data, a tuple of two
states, and weights and biases for dense operations on the
inputs and on the state. It returns a tuple with two members,
an output tensor and a tuple of two new states.
"""
builder = relay.ScopeBuilder()
input_type = relay.TensorType((batch_size, num_hidden), dtype)
weight_type = relay.TensorType((4 * num_hidden, num_hidden), dtype)
bias_type = relay.TensorType((4 * num_hidden,), dtype)
dense_type = relay.TensorType((batch_size, 4 * num_hidden), dtype)
slice_type = relay.TupleType([input_type, input_type, input_type, input_type])
ret_type = relay.TupleType([input_type, relay.TupleType([input_type, input_type])])
inputs = relay.Var("inputs", input_type)
states = relay.Var("states", relay.TupleType([input_type, input_type]))
i2h_weight = relay.Var("i2h_weight", weight_type)
i2h_bias = relay.Var("i2h_bias", bias_type)
h2h_weight = relay.Var("h2h_weight", weight_type)
h2h_bias = relay.Var("h2h_bias", bias_type)
i2h = builder.let(
("i2h", dense_type),
layers.dense_add_bias(
data=inputs, units=num_hidden * 4, weight=i2h_weight, bias=i2h_bias, name="%si2h" % name
),
)
h2h = builder.let(
("h2h", dense_type),
layers.dense_add_bias(
data=relay.TupleGetItem(states, 0),
units=num_hidden * 4,
weight=h2h_weight,
bias=h2h_bias,
name="%sh2h" % name,
),
)
gates = builder.let(("gates", dense_type), relay.add(i2h, h2h))
slice_gates = builder.let(
("slice_gates", slice_type), relay.split(gates, indices_or_sections=4, axis=1).astuple()
)
in_gate = builder.let(
("in_gate", input_type), relay.sigmoid(relay.TupleGetItem(slice_gates, 0))
)
forget_gate = builder.let(
("forget_gate", input_type), relay.sigmoid(relay.TupleGetItem(slice_gates, 1))
)
in_transform = builder.let(
("in_transform", input_type), relay.tanh(relay.TupleGetItem(slice_gates, 2))
)
out_gate = builder.let(
("out_gate", input_type), relay.sigmoid(relay.TupleGetItem(slice_gates, 3))
)
next_c = builder.let(
("next_c", input_type),
relay.add(
relay.multiply(forget_gate, relay.TupleGetItem(states, 1)),
relay.multiply(in_gate, in_transform),
),
)
next_h = builder.let(("next_h", input_type), relay.multiply(out_gate, relay.tanh(next_c)))
ret = builder.let(("ret", ret_type), relay.Tuple([next_h, relay.Tuple([next_h, next_c])]))
builder.ret(ret)
body = builder.get()
return relay.Function(
[inputs, states, i2h_weight, i2h_bias, h2h_weight, h2h_bias], body, ret_type
)
B = T.match_buffer(b, [128, 128], scope="scopeA")
C = T.match_buffer(c, [128, 128], scope="scopeB")
D = T.match_buffer(d, [128, 128], scope="scopeC")
for i, j, k in T.grid(128, 128, 128):
with T.block("update"):
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
with T.init():
D[vi, vj] = C[vi, vj]
D[vi, vj] = D[vi, vj] + A[vi, vk] * B[vj, vk]
gem_ty = relay.FuncType(
[
relay.TupleType([
relay.TensorType((128, 128), "float32"),
relay.TensorType((128, 128), "float32"),
]),
relay.TensorType((128, 128), "float32"),
],
relay.TensorType((128, 128), "float32"),
)
def test_get_prim_func_arg_and_result_constraints():
scopes = tir.analysis.get_prim_func_arg_and_result_memory_constraints(
gem, gem_ty)
assert [x for x in scopes] == ["scopeA", "scopeB", "scopeC"]
def test_apply_prim_func_arg_and_result_memory_constraints():
rewritten = tir.analysis.apply_prim_func_arg_and_result_memory_constraints(
def get_lstm(batch_size, num_hidden, dtype):
'''Returns a module where the main() function is an LSTM RNN,
returning a tuple of two items where the first is the
list of outputs and the second is the final hidden state'''
mod = relay.Module()
p = Prelude(mod)
input_type = relay.TensorType((batch_size, num_hidden), dtype)
weight_type = relay.TensorType((4 * num_hidden, num_hidden), dtype)
bias_type = relay.TensorType((4 * num_hidden, ), dtype)
state_type = relay.TupleType([input_type, input_type])
cell_type = relay.TupleType([input_type, state_type])
state_var_type = relay.TupleType([p.l(input_type), state_type])
input_list = relay.Var('input_list', p.l(input_type))
init_states = relay.Var('init_states', state_type)
cell_fn = lstm_cell(num_hidden, batch_size, dtype, "lstm_cell")
i2h_weight = relay.Var('i2h_weight', weight_type)
i2h_bias = relay.Var('i2h_bias', bias_type)
h2h_weight = relay.Var('h2h_weight', weight_type)
h2h_bias = relay.Var('h2h_bias', bias_type)
state_var = relay.Var('state_var', state_var_type)
input_var = relay.Var('input_var', input_type)
cell_out = relay.Var('cell_out', cell_type)
iteration = relay.Function([state_var, input_var],
relay.Let(
cell_out,
cell_fn(input_var,
relay.TupleGetItem(state_var,
1), i2h_weight,
i2h_bias, h2h_weight, h2h_bias),
relay.Tuple([
p.cons(relay.TupleGetItem(cell_out, 0),
relay.TupleGetItem(state_var,
0)),
relay.TupleGetItem(cell_out, 1)
])), state_var_type)
fold_res = relay.Var('fold_res', state_var_type)
mod['rnn'] = relay.Function(
[i2h_weight, i2h_bias, h2h_weight, h2h_bias, init_states, input_list],
relay.Let(
fold_res,
p.foldl(iteration, relay.Tuple([p.nil(), init_states]),
input_list),
relay.Tuple([
p.rev(relay.TupleGetItem(fold_res, 0)),
relay.TupleGetItem(fold_res, 1)
])), state_var_type)
mod['main'] = relay.Function(
[],
relay.Call(mod.get_global_var('rnn'), [
relay.const(generate_random_tensor(weight_type)),
relay.const(generate_random_tensor(bias_type)),
relay.const(generate_random_tensor(weight_type)),
relay.const(generate_random_tensor(bias_type)),
relay.Tuple([
relay.const(generate_random_tensor(input_type)),
relay.const(generate_random_tensor(input_type))
]),
p.cons(relay.const(generate_random_tensor(input_type)), p.nil())
]))
return mod
