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)