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