def expected():
mod = tvm.IRModule({})
x1 = relay.var("x1", shape=(3, 5))
y1 = relay.var("y1", shape=(3, 5))
sb = relay.ScopeBuilder()
sb.ret(x1 + y1)
fn1 = relay.Function([x1, y1], sb.get())
g1 = relay.GlobalVar("g1")
mod[g1] = fn1
p0 = relay.var("p0", shape=(3, 5))
p1 = relay.var("p1", shape=(3, 5))
p2 = relay.var("p2", shape=(3, 5))
p3 = relay.var("p3", shape=(3, 5))
call_fn2 = p2 - p3
mod["main"] = relay.Function([p0, p1, p2, p3], g1(p0, p1) * call_fn2)
x0 = relay.var("x0", shape=(3, 5))
y0 = relay.var("y0", shape=(3, 5))
z0 = relay.var("z0", shape=(3, 5))
fn0 = relay.Function([x0, y0, z0], x0 - y0 + z0)
g0 = relay.GlobalVar("g0")
mod[g0] = fn0
return mod
def test_custom_op_rel_infer_exception():
"""Tests infer type for custom_op"""
def custom_log1_rel(arg_types, attrs):
assert len(arg_types) == 2, "type relation arg number mismatch!"
return None
op_name = "custom_log2"
_op.register(op_name, r"code(cal log of a tensor.)code")
_op.get(op_name).set_num_inputs(1)
_op.get(op_name).add_argument("data_0", "Tensor", "The input data tensor.")
_op.get(op_name).set_attrs_type_key("DictAttrs")
# call customized relation functions
_op.get(op_name).add_type_rel("custom_log2", custom_log1_rel)
_op.get(op_name).set_support_level(1)
_op.register_pattern(op_name, _op.OpPattern.ELEMWISE)
_op.register_stateful(op_name, False)
def clog(x):
return relay.Call(_op.get(op_name), [x])
tp = relay.TensorType((10, 10), "float32")
x = relay.var("x", tp)
sb = relay.ScopeBuilder()
t1 = sb.let("t1", clog(x))
t2 = sb.let("t2", relay.add(t1, x))
sb.ret(t2)
f = relay.Function([x], sb.get())
with pytest.raises(tvm.error.TVMError) as cm:
fchecked = infer_expr(f)
assert "type relation arg number mismatch" in str(cm.execption)
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 = tvm.IRModule()
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 expected():
sb = relay.ScopeBuilder()
x = relay.var("x", t)
c_folded = (c_data + c_data)
t3 = sb.let("t3", relay.add(relay.const(c_folded), x))
sb.ret(t3)
return relay.Function([x], sb.get())
def get_mod():
mod = tvm.IRModule({})
x = relay.var('x', shape=[], dtype='int32')
fn0 = relay.Function([x], x)
fn0 = fn0.set_attribute("Inline", tvm.tir.IntImm("int32", 1))
gx = relay.GlobalVar("gx")
mod[gx] = fn0
sum_up = relay.GlobalVar('sum_up')
i = relay.var('i', shape=[], dtype='int32')
sb = relay.ScopeBuilder()
with sb.if_scope(relay.equal(i, relay.const(0, dtype="int32"))):
sb.ret(i)
with sb.else_scope():
one_less = relay.subtract(i, relay.const(1, dtype="int32"))
global_call = gx(i)
rec_call = relay.Call(sum_up, [one_less]) + global_call
sb.ret(relay.add(rec_call, i))
func = relay.Function([i],
sb.get(),
ret_type=relay.TensorType([], "int32"))
func = func.set_attribute("Inline", tvm.tir.IntImm("int32", 1))
mod[sum_up] = func
iarg = relay.var("i", shape=[], dtype='int32')
mod["main"] = relay.Function([iarg], sum_up(iarg))
return mod
def test_custom_op_rel_infer():
"""Tests infer type for custom_op"""
def custom_log1_rel(arg_types, attrs):
assert len(arg_types) == 1, "type relation arg number mismatch!"
if attrs:
assert isinstance(attrs, DictAttrs)
inputa_type = arg_types[0]
return relay.TensorType(inputa_type.shape, inputa_type.dtype)
op_name = "custom_log1"
_op.register(op_name, r"code(cal log of a tensor.)code")
_op.get(op_name).set_num_inputs(1)
_op.get(op_name).add_argument("data_0", "Tensor", "The input data tensor.")
_op.get(op_name).set_attrs_type_key("DictAttrs")
# call customized relation functions
_op.get(op_name).add_type_rel("custom_log1", custom_log1_rel)
_op.get(op_name).set_support_level(1)
_op.register_pattern(op_name, _op.OpPattern.ELEMWISE)
_op.register_stateful(op_name, False)
def clog(x):
return relay.Call(_op.get(op_name), [x])
tp = relay.TensorType((10, 10), "float32")
x = relay.var("x", tp)
sb = relay.ScopeBuilder()
t1 = sb.let("t1", clog(x))
t2 = sb.let("t2", relay.add(t1, x))
sb.ret(t2)
f = relay.Function([x], sb.get())
fchecked = infer_expr(f)
assert fchecked.checked_type == relay.FuncType([tp], tp)
def test_recursion():
"""
Program:
def f(n: i32, data: f32) -> f32 {
if (n == 0) {
return data;
} else {
return 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 "%3 = @f(%1, %2)" in mod.astext()
assert mod[f].checked_type == relay.FuncType([ti32, tf32], tf32)
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_recursive_func():
mod = tvm.IRModule({})
x = relay.var('x', shape=[], dtype='int32')
fn0 = relay.Function([x], x)
gx = relay.GlobalVar("gx")
mod[gx] = fn0
sum_up = relay.GlobalVar('sum_up')
i = relay.var('i', shape=[], dtype='int32')
sb = relay.ScopeBuilder()
with sb.if_scope(relay.equal(i, relay.const(0, dtype='int32'))):
sb.ret(i)
with sb.else_scope():
one_less = relay.subtract(i, relay.const(1, dtype='int32'))
global_call = gx(i)
rec_call = relay.Call(sum_up, [one_less]) + global_call
sb.ret(relay.add(rec_call, i))
func = relay.Function([i],
sb.get(),
ret_type=relay.TensorType([], 'int32'))
func = func.with_attr("Compiler", "a")
mod[sum_up] = func
iarg = relay.var('i', shape=[], dtype='int32')
mod["main"] = relay.Function([iarg], sum_up(iarg))
call_graph = relay.analysis.CallGraph(mod)
assert call_graph.is_recursive(sum_up)
assert call_graph.ref_count(sum_up) == 2
assert call_graph.ref_count(gx) == 1
assert call_graph.ref_count("main") == 0
def test_checkpoint(executor_kind, target, dev):
inputs = [relay.var("x{}".format(i), shape=(1, )) for i in range(4)]
output = relay.multiply(relay.add(inputs[0], inputs[1]),
relay.add(inputs[2], inputs[3]))
check_grad(relay.Function(inputs, relay.annotation.checkpoint(output)),
executor_kind=executor_kind)
scope = relay.ScopeBuilder()
out_tuple = scope.let(
"out_tuple",
relay.Tuple([
relay.add(inputs[0], inputs[1]),
relay.multiply(inputs[2], inputs[3])
]),
)
scope.ret(
relay.subtract(
relay.annotation.checkpoint(relay.TupleGetItem(out_tuple, 0)),
relay.TupleGetItem(out_tuple, 1),
))
out_single = scope.get()
check_grad(
relay.Function(inputs, out_single),
target_devices=[(target, dev)],
executor_kind=executor_kind,
)
def __init__(self, params, quantized_dtypes):
ExprMutator.__init__(self)
self.params = set(params)
self.quantized_dtypes = quantized_dtypes
self.subtree_params = set()
self.new_func_params = []
self.prefix_sb = relay.ScopeBuilder()
self.prefix_binding_map = {}
def before():
sb = relay.ScopeBuilder()
x = relay.var("x", t)
t1 = sb.let("t1", relay.const(c_data))
t2 = sb.let("t2", relay.add(t1, t1))
t3 = sb.let("t3", relay.add(t2, x))
sb.ret(t3)
return relay.Function([x], sb.get())
def before(dshape, dtype):
x = relay.var('x', shape=dshape, dtype=dtype)
y = relay.var('y', shape=dshape, dtype=dtype)
sb = relay.ScopeBuilder()
with sb.if_scope(relay.op.greater(x, y)):
sb.ret(relay.Function([], x))
with sb.else_scope():
sb.ret(relay.Function([], y))
return relay.Function([x, y], relay.Call(sb.get(), []))
def expected():
sb = relay.ScopeBuilder()
x = relay.var("x", t)
c_folded = c_data + c_data
t3 = sb.let(
"t3", annot_expr(relay.add(annot_expr(relay.const(c_folded)), x)))
sb.ret(t3)
f = relay.Function([x], sb.get())
return annot_func(f)
def test_monomorphic_let():
"Program: let %x = 1; %x"
# TODO(@jroesch): this seems whack.
sb = relay.ScopeBuilder()
x = relay.var("x", dtype="float64", shape=())
x = sb.let(x, relay.const(1.0, "float64"))
sb.ret(x)
xchecked = infer_expr(sb.get())
assert xchecked.checked_type == relay.scalar_type("float64")
def func1():
sb = relay.ScopeBuilder()
p0 = relay.var("p0", shape=shape)
p1 = relay.var("p1", shape=shape)
a0 = sb.let("a0", relay.add(p0, relay.const(1)))
a1 = sb.let("a1", relay.add(p1, relay.const(1)))
a2 = sb.let("a2", relay.add(a0, a1))
sb.ret(a2)
return relay.Function([p0, p1], sb.get())
def func2():
# Alpha conversion is structurally equal
sb = relay.ScopeBuilder()
p0 = relay.var("p0", shape=shape)
p1 = relay.var("p1", shape=shape)
a1 = sb.let("a1", relay.add(p0, relay.const(1)))
a0 = sb.let("a0", relay.add(p1, relay.const(1)))
a2 = sb.let("a2", relay.add(a1, a0))
sb.ret(a2)
return relay.Function([p0, p1], sb.get())
def func3():
# But changing the order of bindings is not structurally equal
# (even though algebraically equal)
sb = relay.ScopeBuilder()
p0 = relay.var("p0", shape=shape)
p1 = relay.var("p1", shape=shape)
a1 = sb.let("a1", relay.add(p1, relay.const(1)))
a0 = sb.let("a0", relay.add(p0, relay.const(1)))
a2 = sb.let("a2", relay.add(a1, a0))
sb.ret(a2)
return relay.Function([p0, p1], sb.get())
def before():
sb = relay.ScopeBuilder()
x = relay.var("x", t)
x.virtual_device_ = tvm.cpu()
t1 = sb.let("t1", annot_expr(relay.const(c_data)))
t2 = sb.let("t2", annot_expr(relay.add(t1, t1)))
t3 = sb.let("t3", annot_expr(relay.add(t2, x)))
sb.ret(t3)
f = relay.Function([x], sb.get())
f.virtual_device_ = tvm.cpu()
return f
def expected():
sb = relay.ScopeBuilder()
x = relay.var("x", t)
x.virtual_device_ = tvm.cpu()
c_folded = c_data + c_data
t3 = sb.let(
"t3", annot_expr(relay.add(annot_expr(relay.const(c_folded)), x)))
sb.ret(t3)
f = relay.Function([x], sb.get())
f.virtual_device_ = tvm.cpu()
return f
def build_relay_module(batch_size, input_size, hidden_size, time_steps,
dense_dim):
mod = tvm.IRModule()
mod["lstm_layer"] = lstm_definition(batch_size, input_size, hidden_size,
time_steps)
mod["linear_layer"] = linear_layer_definition(batch_size, hidden_size,
dense_dim)
lstm_var = mod.get_global_var("lstm_layer")
linear_var = mod.get_global_var("linear_layer")
# now we build up our main function
input_var = relay.var("input", shape=(batch_size, time_steps, input_size))
init_hidden_var = relay.var("init_hidden", shape=(batch_size, hidden_size))
init_cell_var = relay.var("init_cell", shape=(batch_size, hidden_size))
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))
lstm_bias_var = relay.var("lstm_bias", shape=(4 * hidden_size, ))
linear_weight_var = relay.var("linear_weight",
shape=(dense_dim, hidden_size))
linear_bias_var = relay.var("linear_bias", shape=(dense_dim, ))
builder = relay.ScopeBuilder()
state_var = builder.let("state",
relay.Tuple([init_hidden_var, init_cell_var]))
lstm_res = builder.let(
"lstm_res",
lstm_var(
input_var,
state_var,
i2h_weight_var,
h2h_weight_var,
lstm_bias_var,
# the keras model only gave one bias,
# so set the other to zero
# (hopefully this is correct)
relay.zeros_like(lstm_bias_var)))
final_hidden = builder.let("final_hidden", relay.TupleGetItem(lstm_res, 1))
# to match PT's semantics, we're undoing the reshape in LSTM :)
reshape_hidden = builder.let("reshape_hidden",
relay.squeeze(final_hidden, axis=[0]))
linear_result = builder.let(
"linear_result",
linear_var(reshape_hidden, linear_weight_var, linear_bias_var))
# finally do a softmax
builder.ret(relay.nn.softmax(linear_result))
main_func = relay.Function([
input_var, init_hidden_var, init_cell_var, i2h_weight_var,
h2h_weight_var, lstm_bias_var, linear_weight_var, linear_bias_var
], builder.get())
mod["main"] = main_func
return mod
def test_decl():
"""Program:
def f(x : Tensor[(10, 10), f32]) {
return log(x);
}
"""
sb = relay.ScopeBuilder()
tp = relay.TensorType((10, 10))
x = relay.var("x", tp)
f = relay.Function([x], relay.log(x))
fchecked = relay.ir_pass.infer_type(f)
assert fchecked.checked_type == relay.FuncType([tp], tp)
def expected(dshape, dtype):
x = relay.var('x', shape=dshape, dtype=dtype)
y = relay.var('y', shape=dshape, dtype=dtype)
sb = relay.ScopeBuilder()
p1 = relay.var('p1', shape=dshape, dtype=dtype)
p2 = relay.var('p2', shape=dshape, dtype=dtype)
fused_gt = relay.Function([p1, p2], relay.op.greater(p1, p2))
with sb.if_scope(fused_gt(x, y)):
sb.ret(relay.Function([], x))
with sb.else_scope():
sb.ret(relay.Function([], y))
return relay.Function([x, y], relay.Call(sb.get(), []))
def test_compose():
mod = relay.Module()
p = Prelude(mod)
compose = p.compose
# remove all functions to not have pattern match to pass vm compilation
# TODO(wweic): remove the hack and implement pattern match
for v, _ in mod.functions.items():
if v.name_hint == 'compose':
continue
mod[v] = relay.const(0)
# add_one = fun x -> x + 1
sb = relay.ScopeBuilder()
x = relay.var('x', 'float32')
x1 = sb.let('x1', x)
xplusone = x1 + relay.const(1.0, 'float32')
sb.ret(xplusone)
body = sb.get()
add_one = relay.GlobalVar("add_one")
add_one_func = relay.Function([x], body)
# add_two = compose(add_one, add_one)
sb = relay.ScopeBuilder()
y = relay.var('y', 'float32')
add_two_func = sb.let('add_two', compose(add_one_func, add_one_func))
add_two_res = add_two_func(y)
sb.ret(add_two_res)
add_two_body = sb.get()
mod[add_one] = add_one_func
f = relay.Function([y], add_two_body)
mod["main"] = f
x_data = np.array(np.random.rand()).astype('float32')
result = veval(mod)(x_data)
tvm.testing.assert_allclose(result.asnumpy(), x_data + 2.0)
def test_adt_compose():
mod = relay.Module()
p = Prelude(mod)
compose = p.compose
# add_one = fun x -> x + 1
sb = relay.ScopeBuilder()
x = relay.var('x', 'float32')
x1 = sb.let('x1', x)
xplusone = x1 + relay.const(1.0, 'float32')
sb.ret(xplusone)
body = sb.get()
add_one = relay.GlobalVar("add_one")
add_one_func = relay.Function([x], body)
# add_two = compose(add_one, add_one)
sb = relay.ScopeBuilder()
y = relay.var('y', 'float32')
add_two_func = sb.let('add_two', compose(add_one_func, add_one_func))
add_two_res = add_two_func(y)
sb.ret(add_two_res)
add_two_body = sb.get()
mod[add_one] = add_one_func
f = relay.Function([y], add_two_body)
mod["main"] = f
vm = create_vm(mod)
ser = serializer.Serializer(vm)
code, lib = ser.serialize()
deser = deserializer.Deserializer(code, lib)
des_vm = deser.deserialize()
x_data = np.array(np.random.rand()).astype('float32')
result = veval(des_vm, x_data)
tvm.testing.assert_allclose(result.asnumpy(), x_data + 2.0)
def expected(dshape, dtype):
x = relay.var("x", shape=dshape, dtype=dtype)
y = relay.var("y", shape=dshape, dtype=dtype)
sb = relay.ScopeBuilder()
p1 = relay.var("p1", shape=dshape, dtype=dtype)
p2 = relay.var("p2", shape=dshape, dtype=dtype)
fused_gt = relay.Function([p1, p2], relay.op.greater(p1, p2))
fused_gt = fused_gt.with_attr("Primitive", tvm.tir.IntImm("int32", 1))
with sb.if_scope(fused_gt(x, y)):
sb.ret(relay.Function([], x))
with sb.else_scope():
sb.ret(relay.Function([], y))
return relay.Function([x, y], relay.Call(sb.get(), []))
def test_adt_compose():
mod = relay.Module()
p = Prelude(mod)
compose = p.compose
# add_one = fun x -> x + 1
sb = relay.ScopeBuilder()
x = relay.var('x', 'float32')
x1 = sb.let('x1', x)
xplusone = x1 + relay.const(1.0, 'float32')
sb.ret(xplusone)
body = sb.get()
add_one = relay.GlobalVar("add_one")
add_one_func = relay.Function([x], body)
# add_two = compose(add_one, add_one)
sb = relay.ScopeBuilder()
y = relay.var('y', 'float32')
add_two_func = sb.let('add_two', compose(add_one_func, add_one_func))
add_two_res = add_two_func(y)
sb.ret(add_two_res)
add_two_body = sb.get()
mod[add_one] = add_one_func
f = relay.Function([y], add_two_body)
mod["main"] = f
exe = create_exec(mod)
code, lib = exe.save()
des_exec = _vm.Executable.load_exec(code, lib)
des_vm = _vm.VirtualMachine(des_exec)
des_vm.init(tvm.cpu())
x_data = np.array(np.random.rand()).astype('float32')
result = veval(des_vm, x_data)
tvm.testing.assert_allclose(result.asnumpy(), x_data + 2.0)
def fuse_partitions(pre_mod, mid_mod, post_mod):
"""Combine prefix, middle, and suffix modules into a single module.
The combined module includes an additional `main` that fuses all three
partitions together.
Parameters
----------
pre_mod : tvm.IRModule
Module containing an input quantization function
mid_mod : tvm.IRModule
Module containing core of a quantized inference function
post_mod : tvm.IRModule
Module containing an output dequantization function
Returns
-------
fused_mod : tvm.IRModule
Module containing the input quantization, core quantized inference,
output dequantization, and full quantized inference functions
"""
pre_func = pre_mod['main']
mid_func = mid_mod['main']
post_func = post_mod['main']
# create a module containing the prefix, middle, and suffix partitions
fused_mod = tvm.IRModule(functions={
relay.GlobalVar('quantize_inputs'): pre_func,
relay.GlobalVar('quantized_main'): mid_func,
relay.GlobalVar('dequantize_outputs'): post_func,
})
# construct a `main` that strings together the partitions, such that its
# behaviour is equivalent to `main` in an *unpartitioned* module
scope_builder = relay.ScopeBuilder()
fused_mod_main_params = [relay.Var(param.name_hint) for param in pre_func.params]
quantized_inputs = scope_builder.let('quantized_inputs', relay.Call(
fused_mod.get_global_var('quantize_inputs'),
fused_mod_main_params
))
quantized_outputs = scope_builder.let('quantized_outputs', relay.Call(
fused_mod.get_global_var('quantized_main'),
[relay.TupleGetItem(quantized_inputs, i) for i in range(len(pre_func.ret_type.fields))]
))
dequantized_outputs = scope_builder.let('dequantized_outputs', relay.Call(
fused_mod.get_global_var('dequantize_outputs'),
[quantized_outputs]
))
scope_builder.ret(dequantized_outputs)
fused_mod['main'] = relay.Function(fused_mod_main_params, scope_builder.get())
return fused_mod
def test_let_scalar():
sb = relay.ScopeBuilder()
x = relay.var('x', 'float32')
x1 = sb.let('x1', x)
xplusone = x1 + relay.const(42.0, 'float32')
sb.ret(xplusone)
body = sb.get()
f = relay.Function([x], body)
x_data = np.array(np.random.rand()).astype('float32')
result = veval(f, x_data)
tvm.testing.assert_allclose(result.asnumpy(), x_data + 42.0)
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)
