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