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 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)
