def run(dtype):
mod = tvm.IRModule()
p = Prelude(mod)
l = relay.var("l")
i = relay.var("i")
read_func = p.get_global_var("tensor_array_read", dtype)
tensor_array = p.get_global_var("tensor_array", dtype)
mod["main"] = relay.Function([l, i], read_func(tensor_array(l), i))
expected = [0]
check_tensor_array(mod, expected, *(1, 0), dtype=dtype)
check_tensor_array(mod, expected, *(5, 1), dtype=dtype)
def run(dtype):
mod = tvm.IRModule()
p = Prelude(mod)
# tensor array
v1 = relay.var("v1")
v2 = relay.var("v2")
v3 = relay.var("v2")
tensor_array = p.get_global_var("tensor_array", dtype)
tensor_array1 = tensor_array(relay.const(3))
write_func = p.get_global_var("tensor_array_write", dtype)
scatter_func = p.get_global_var("tensor_array_scatter", dtype)
tensor2 = p.get_tensor_ctor("tensor2", dtype)
tensor_array1 = write_func(tensor_array1, relay.const(0), tensor2(v1))
tensor_array1 = write_func(tensor_array1, relay.const(1), tensor2(v2))
tensor_array1 = write_func(tensor_array1, relay.const(2), tensor2(v3))
# indices array
index = relay.var("index")
# values array
value_0 = relay.var("value_0")
value_1 = relay.var("value_1")
values_array = tensor_array(relay.const(2))
values_array = write_func(values_array, relay.const(0),
tensor2(value_0))
values_array = write_func(values_array, relay.const(1),
tensor2(value_1))
# create the scatter function
tensor_array_scatter = scatter_func(tensor_array1, index, values_array)
mod["main"] = relay.Function([v1, v2, v3, index, value_0, value_1],
tensor_array_scatter)
# initialize and check
v1_data = np.random.uniform(low=0.0, high=8.0,
size=(2, 3)).astype(dtype)
v2_data = np.random.uniform(low=0.0, high=8.0,
size=(2, 3)).astype(dtype)
v3_data = np.random.uniform(low=0.0, high=8.0,
size=(2, 3)).astype(dtype)
index_data = np.array([0, 1], dtype="int32")
val1_data = np.random.uniform(low=0.0, high=8.0,
size=(2, 3)).astype(dtype)
val2_data = np.random.uniform(low=0.0, high=8.0,
size=(2, 3)).astype(dtype)
expected = [val1_data, val2_data, v3_data]
check_tensor_array(
mod,
expected,
*(v1_data, v2_data, v3_data, index_data, val1_data, val2_data),
dtype=dtype,
)
def run(dtype):
mod = tvm.IRModule()
p = Prelude(mod)
tensor_t = p.get_type("tensor_t", dtype)
v1 = relay.var("v1")
v2 = relay.var("v2")
tensor_array = p.get_global_var("tensor_array", dtype)
init_tensor_array = tensor_array(relay.const(2))
write_func = p.get_global_var("tensor_array_write", dtype)
tensor1 = p.get_tensor_ctor("tensor1", dtype)
tensor_array1 = write_func(init_tensor_array, relay.const(0), tensor1(v1))
tensor_array2 = write_func(tensor_array1, relay.const(1), tensor1(v2))
mod["main"] = relay.Function([v1, v2], tensor_array2)
expected = [3, 7]
check_tensor_array(mod, expected, *(3, 7), dtype=dtype)
def test_globalvar_as_call_arg():
mod = tvm.IRModule()
p = Prelude(mod)
tensor_array = p.get_global_var("tensor_array", "int32")
tensor1 = p.get_ctor(p.get_name("tensor_t", "int32"), "tensor1", "int32")
write = p.get_global_var("tensor_array_write", "int32")
stack = p.get_global_var("tensor_array_stack", "int32")
v = relay.var("v")
init_tensor_array = tensor_array(relay.const(3))
tensor_array1 = write(init_tensor_array, relay.const(0), tensor1(v))
tensor_array2 = stack(tensor_array1)
mod["main"] = relay.Function([v], tensor_array2)
mod = relay.transform.RemoveUnusedFunctions()(mod)
l = set([x[0].name_hint for x in mod.functions.items()])
assert "tensor_array_int32" in l
def run(dtype):
x = relay.var("x")
mod = tvm.IRModule()
p = Prelude(mod)
tensor_array = p.get_global_var("tensor_array", dtype)
mod["main"] = relay.Function([x], tensor_array(x))
expected = np.array([0, 0, 0, 0, 0])
check_tensor_array(mod, expected, 5, dtype=dtype)
def run(dtype):
mod = tvm.IRModule()
p = Prelude(mod)
unstack_tensor1 = p.get_global_var("tensor_array_unstack_tensor1", dtype)
v = relay.var("v")
mod["main"] = relay.Function([v], unstack_tensor1(v))
t = np.random.uniform(low=0.0, high=8.0, size=(1,)).astype(dtype)
check_tensor_array(mod, t, t, dtype=dtype)
def run(dtype):
mod = tvm.IRModule()
p = Prelude(mod)
# tensor array
v1 = relay.var("v1")
v2 = relay.var("v2")
v3 = relay.var("v2")
tensor_array = p.get_global_var("tensor_array", dtype)
tensor_array1 = tensor_array(relay.const(3))
write_func = p.get_global_var("tensor_array_write", dtype)
split_func = p.get_global_var("tensor_array_split", dtype)
tensor2 = p.get_tensor_ctor("tensor2", dtype)
tensor_array1 = write_func(tensor_array1, relay.const(0), tensor2(v1))
tensor_array1 = write_func(tensor_array1, relay.const(1), tensor2(v2))
tensor_array1 = write_func(tensor_array1, relay.const(2), tensor2(v3))
# value tensor
value = relay.var("value")
# lengths tensor
ta_len = relay.var("length")
# create the scatter function
tensor_array_split = split_func(tensor_array1, tensor2(value), ta_len)
mod["main"] = relay.Function([v1, v2, v3, value, ta_len],
tensor_array_split)
# initialize and check
v1_data = np.random.uniform(low=0.0, high=8.0,
size=(2, 3)).astype(dtype)
v2_data = np.random.uniform(low=0.0, high=8.0,
size=(2, 3)).astype(dtype)
v3_data = np.random.uniform(low=0.0, high=8.0,
size=(2, 3)).astype(dtype)
value_data = np.random.uniform(low=0.0, high=8.0,
size=(4, 3)).astype(dtype)
length_data = np.array([2, 2], dtype="int32")
expected = np.concatenate([value_data, v3_data])
expected = np.split(expected, indices_or_sections=[2, 4])
check_tensor_array(mod,
expected,
*(v1_data, v2_data, v3_data, value_data,
length_data),
dtype=dtype)
def run(dtype):
mod = tvm.IRModule()
p = Prelude(mod)
v1 = relay.var("v1")
v2 = relay.var("v2")
tensor_array = p.get_global_var("tensor_array", dtype)
tensor_array1 = tensor_array(relay.const(2))
write_func = p.get_global_var("tensor_array_write", dtype)
concat_func = p.get_global_var("tensor_array_concat", dtype)
tensor1 = p.get_tensor_ctor("tensor2", dtype)
tensor_array1 = write_func(tensor_array1, relay.const(0), tensor1(v1))
tensor_array1 = write_func(tensor_array1, relay.const(1), tensor1(v2))
tensor_array_concat = concat_func(tensor_array1)
mod["main"] = relay.Function([v1, v2], tensor_array_concat)
v1_data = np.random.uniform(low=0.0, high=8.0, size=(2, 3)).astype(dtype)
v2_data = np.random.uniform(low=0.0, high=8.0, size=(1, 3)).astype(dtype)
expected = [np.concatenate((v1_data, v2_data), axis=0)]
check_tensor_array(mod, expected, *(v1_data, v2_data), dtype=dtype)
def run(dtype):
x = relay.var("x")
mod = tvm.IRModule()
p = Prelude(mod)
expand_dims_func = p.get_global_var("tensor_expand_dims", dtype)
tensor1 = p.get_tensor_ctor("tensor1", dtype)
mod["main"] = relay.Function([x], expand_dims_func(tensor1(x)))
x_np = np.random.uniform(low=0.0, high=8.0, size=(1, )).astype(dtype)
expected = [np.expand_dims(x_np, axis=0)]
check_tensor_array(mod, expected, x_np)
def run(dtype):
mod = tvm.IRModule()
p = Prelude(mod)
tensor_t = p.get_type("tensor_t", dtype)
rlist = p.mod.get_global_type_var(f"List")
tensor_array = p.get_global_var("tensor_array", dtype)
tensor1 = p.get_tensor_ctor("tensor1", dtype)
write = p.get_global_var("tensor_array_write", dtype)
stack = p.get_global_var("tensor_array_stack", dtype)
# TODO extract test case from inference failures
# setting this wrong causes crashes
v = relay.var("v", shape=(1, ), dtype=dtype)
init_tensor_array = tensor_array(relay.const(3))
tensor_array1 = write(init_tensor_array, relay.const(0), tensor1(v))
tensor_array2 = write(tensor_array1, relay.const(1), tensor1(v))
tensor_array3 = write(tensor_array2, relay.const(2), tensor1(v))
tensor_array4 = stack(tensor_array3)
mod["main"] = relay.Function([v], tensor_array4, tensor_t())
t = np.random.uniform(low=0.0, high=8.0, size=(1, )).astype(dtype)
expected = [np.stack([t, t, t])]
check_tensor_array(mod, expected, t, dtype=dtype)
def run(dtype):
mod = tvm.IRModule()
p = Prelude(mod)
concat = p.get_global_var("tensor_concatenate", dtype)
tensor1 = p.get_tensor_ctor("tensor1", dtype)
v1 = relay.var("v1", shape=(tvm.tir.Any(),), dtype=dtype)
v2 = relay.var("v2", shape=(tvm.tir.Any(),), dtype=dtype)
mod["main"] = relay.Function([v1, v2], concat(tensor1(v1), tensor1(v2)))
v1_data = np.random.uniform(low=0.0, high=8.0, size=(5,)).astype(dtype)
v2_data = np.random.uniform(low=0.0, high=8.0, size=(5,)).astype(dtype)
expected = [np.concatenate((v1_data, v2_data))]
check_tensor_array(mod, expected, *(v1_data, v2_data), dtype=dtype)
def run(dtype):
mod = tvm.IRModule()
p = Prelude(mod)
take = p.get_global_var("tensor_take", dtype)
tensor2 = p.get_tensor_ctor("tensor2", dtype)
v = relay.var("v")
lower = relay.var("lower")
upper = relay.var("upper")
mod["main"] = relay.Function([v, lower, upper], take(tensor2(v), lower, upper))
v_data = np.random.uniform(low=0.0, high=8.0, size=(10, 10)).astype(dtype)
expected = [np.take(v_data, range(2, 5), axis=0)]
check_tensor_array(mod, expected, *(v_data, 2, 5), dtype=dtype)
expected = [np.take(v_data, range(0, 9), axis=0)]
check_tensor_array(mod, expected, *(v_data, 0, 9), dtype=dtype)
