def run(dtype, shape):
mod = tvm.IRModule()
p = Prelude(mod)
static_tensor_array_ops = StaticTensorArrayOps(p, dtype, shape)
static_tensor_array_ops.register()
v1 = relay.var("v1")
v2 = relay.var("v2")
tensor_array = p.get_global_var_static("tensor_array", dtype, shape)
tensor_array1 = tensor_array(relay.const(2))
write_func = p.get_global_var_static("tensor_array_write", dtype,
shape)
concat_func = p.get_global_var_static("tensor_array_concat", dtype,
shape)
tensor = p.get_tensor_ctor_static("tensor_constructor", dtype, shape)
tensor_array1 = write_func(tensor_array1, relay.const(0), tensor(v1))
tensor_array1 = write_func(tensor_array1, relay.const(1), tensor(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, shape):
mod = tvm.IRModule()
p = Prelude(mod)
static_tensor_array_ops = StaticTensorArrayOps(p, dtype, shape)
static_tensor_array_ops.register()
ta_length = 2
np_data_list = [
np.random.uniform(0, 10, size=shape).astype(dtype)
for _ in range(ta_length)
]
v0 = relay.var("v0")
v1 = relay.var("v1")
tensor_array = p.get_global_var_static("tensor_array", dtype, shape)
init_tensor_array = tensor_array(relay.const(ta_length))
write_func = p.get_global_var_static("tensor_array_write", dtype,
shape)
tensor = p.get_tensor_ctor_static("tensor_constructor", dtype, shape)
tensor_array0 = write_func(init_tensor_array, relay.const(0),
tensor(v0))
tensor_array1 = write_func(tensor_array0, relay.const(1), tensor(v1))
mod["main"] = relay.Function([v0, v1], tensor_array1)
expected = np_data_list
check_tensor_array(mod, expected, *np_data_list, dtype=dtype)
def run(dtype, shape, indices_shape=None):
mod = tvm.IRModule()
p = Prelude(mod)
static_tensor_array_ops = StaticTensorArrayOps(p, dtype, shape)
static_tensor_array_ops.register()
if indices_shape is not None:
static_tensor_array_ops.define_tensor_array_scatter(
indices_shape, True)
# tensor array
v1 = relay.var("v1")
v2 = relay.var("v2")
v3 = relay.var("v2")
tensor_array = p.get_global_var_static("tensor_array", dtype, shape)
tensor_array0 = tensor_array(relay.const(3))
write_func = p.get_global_var_static("tensor_array_write", dtype,
shape)
scatter_func = p.get_global_var_static("tensor_array_scatter", dtype,
shape)
tensor = p.get_tensor_ctor_static("tensor_constructor", dtype, shape)
tensor_array1 = write_func(tensor_array0, relay.const(0), tensor(v1))
tensor_array1 = write_func(tensor_array1, relay.const(1), tensor(v2))
tensor_array1 = write_func(tensor_array1, relay.const(2), tensor(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),
tensor(value_0))
values_array = write_func(values_array, relay.const(1),
tensor(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=shape).astype(dtype)
v2_data = np.random.uniform(low=0.0, high=8.0,
size=shape).astype(dtype)
v3_data = np.random.uniform(low=0.0, high=8.0,
size=shape).astype(dtype)
index_data = np.array([0, 1], dtype="int32")
val1_data = np.random.uniform(low=0.0, high=8.0,
size=shape).astype(dtype)
val2_data = np.random.uniform(low=0.0, high=8.0,
size=shape).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, shape):
mod = tvm.IRModule()
p = Prelude(mod)
static_tensor_array_ops = StaticTensorArrayOps(p, dtype, shape)
static_tensor_array_ops.register()
np_data_list = []
ta_length = 3
for _ in range(ta_length):
np_data_list.append(np.random.uniform(0, 10, size=shape).astype(dtype))
v0 = relay.var("v0")
v1 = relay.var("v1")
v2 = relay.var("v2")
n = relay.var("n")
tensor = p.get_tensor_ctor_static("tensor_constructor", dtype, shape)
tensor_array = p.get_global_var_static("tensor_array", dtype, shape)
init_tensor_array = tensor_array(relay.const(ta_length))
read_func = p.get_global_var_static("tensor_array_read", dtype, shape)
write_func = p.get_global_var_static("tensor_array_write", dtype, shape)
get_data_func = p.get_global_var_static("tensor_get_data", dtype, shape)
tensor_array0 = write_func(init_tensor_array, relay.const(0), tensor(v0))
tensor_array1 = write_func(tensor_array0, relay.const(1), tensor(v1))
tensor_array2 = write_func(tensor_array1, relay.const(2), tensor(v2))
mod["main"] = relay.Function([v0, v1, v2, n], get_data_func(read_func(tensor_array2, n)))
expected = [np_data_list[0]]
check_tensor_array(mod, expected, *list(np_data_list + [0]), dtype=dtype)
expected = [np_data_list[1]]
check_tensor_array(mod, expected, *list(np_data_list + [1]), dtype=dtype)
expected = [np_data_list[2]]
check_tensor_array(mod, expected, *list(np_data_list + [2]), dtype=dtype)
def run(dtype, shape):
mod = tvm.IRModule()
p = Prelude(mod)
static_tensor_array_ops = StaticTensorArrayOps(p, dtype, shape)
static_tensor_array_ops.register()
tensor_array = p.get_global_var_static("tensor_array", dtype, shape)
tensor = p.get_tensor_ctor_static("tensor_constructor", dtype, shape)
write = p.get_global_var_static("tensor_array_write", dtype, shape)
stack = p.get_global_var_static("tensor_array_stack", dtype, shape)
v = relay.var("v")
init_tensor_array = tensor_array(relay.const(3))
tensor_array1 = write(init_tensor_array, relay.const(0), tensor(v))
tensor_array2 = write(tensor_array1, relay.const(1), tensor(v))
tensor_array3 = write(tensor_array2, relay.const(2), tensor(v))
tensor_array4 = stack(tensor_array3)
mod["main"] = relay.Function([v], tensor_array4)
t = np.random.uniform(low=0.0, high=8.0, size=shape).astype(dtype)
expected = [np.stack([t, t, t])]
check_tensor_array(mod, expected, t, dtype=dtype)
def run(dtype, shape):
mod = tvm.IRModule()
p = Prelude(mod)
static_tensor_array_ops = StaticTensorArrayOps(p, dtype, shape)
static_tensor_array_ops.register()
unstack_tensor = p.get_global_var_static("tensor_array_unstack", dtype, shape)
v = relay.var("v")
mod["main"] = relay.Function([v], unstack_tensor(v))
t = np.random.uniform(low=0, high=10, size=shape).astype(dtype)
(*expected,) = t
check_tensor_array(mod, expected, t, dtype=dtype)
def run(dtype, shape):
mod = tvm.IRModule()
p = Prelude(mod)
static_tensor_array_ops = StaticTensorArrayOps(p, dtype, shape)
static_tensor_array_ops.register()
tensor_array = p.get_global_var_static("tensor_array", dtype, shape)
tensor = p.get_tensor_ctor_static("tensor_constructor", dtype, shape)
write = p.get_global_var_static("tensor_array_write", dtype, shape)
gather = p.get_global_var_static("tensor_array_gather", dtype, shape)
v = relay.var("v")
indice = relay.var("indice")
init_tensor_array = tensor_array(relay.const(3))
tensor_array1 = write(init_tensor_array, relay.const(0), tensor(v))
tensor_array2 = write(tensor_array1, relay.const(1), tensor(v))
tensor_array3 = write(tensor_array2, relay.const(2), tensor(v))
out = gather(tensor_array3, indice)
mod["main"] = relay.Function([v, indice], out)
t = np.random.uniform(low=0.0, high=8.0, size=shape).astype(dtype)
indice_data = np.array([0, 2], dtype="int32")
expected = [np.stack([t, t])]
check_tensor_array(mod, expected, *(t, indice_data), dtype=dtype)
def run(dtype, shape):
x = relay.var("x")
mod = tvm.IRModule()
p = Prelude(mod)
static_tensor_array_ops = StaticTensorArrayOps(p, dtype, shape)
static_tensor_array_ops.register()
expand_dims_func = p.get_global_var_static("tensor_expand_dims", dtype, shape)
tensor = p.get_tensor_ctor_static("tensor_constructor", dtype, shape)
mod["main"] = relay.Function([x], expand_dims_func(tensor(x)))
x_np = np.random.uniform(low=0.0, high=8.0, size=shape).astype(dtype)
expected = [np.expand_dims(x_np, axis=0)]
check_tensor_array(mod, expected, x_np)
def run(dtype, shape):
mod = tvm.IRModule()
p = Prelude(mod)
static_tensor_array_ops = StaticTensorArrayOps(p, dtype, shape)
static_tensor_array_ops.register()
concat = p.get_global_var_static("tensor_concatenate", dtype, shape)
tensor = p.get_tensor_ctor_static("tensor_constructor", dtype, shape)
v1 = relay.var("v1")
v2 = relay.var("v2")
mod["main"] = relay.Function([v1, v2], concat(tensor(v1), tensor(v2)))
v1_data = np.random.uniform(low=0.0, high=8.0, size=shape).astype(dtype)
v2_data = np.random.uniform(low=0.0, high=8.0, size=shape).astype(dtype)
expected = [np.concatenate((v1_data, v2_data))]
check_tensor_array(mod, expected, *(v1_data, v2_data), dtype=dtype)
def run(dtype, shape):
mod = tvm.IRModule()
p = Prelude(mod)
static_tensor_array_ops = StaticTensorArrayOps(p, dtype, shape)
static_tensor_array_ops.register()
take = p.get_global_var_static("tensor_take", dtype, shape)
tensor_constructor = p.get_tensor_ctor_static("tensor_constructor", dtype, shape)
v = relay.var("v")
lower = relay.var("lower")
upper = relay.var("upper")
mod["main"] = relay.Function([v, lower, upper], take(tensor_constructor(v), lower, upper))
v_data = np.random.uniform(low=0.0, high=8.0, size=shape).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)
