| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, |
| # software distributed under the License is distributed on an |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| # KIND, either express or implied. See the License for the |
| # specific language governing permissions and limitations |
| # under the License. |
| import pytest |
| import tvm |
| from tvm import relay |
| from tvm.relay import testing |
| from tvm.relay.backend.interpreter import ConstructorValue |
| from tvm.relay import create_executor |
| from tvm.relay.prelude import Prelude, StaticTensorArrayOps |
| from tvm.relay.testing import count as count_, make_nat_value, make_nat_expr |
| |
| import numpy as np |
| |
| |
| def vmobj_to_list(mod, o, dtype="float32"): |
| _, tensor_nil, _, _, _, _, _, _, _ = mod.get_type(f"tensor_{dtype}_t") |
| if isinstance(o, tvm.nd.NDArray): |
| return [o.numpy().tolist()] |
| elif isinstance(o, tvm.runtime.container.ADT): |
| if len(o) == 0: |
| if tensor_nil.tag == o.tag: |
| return [0] |
| return [] |
| |
| result = [] |
| for f in o: |
| result.extend(vmobj_to_list(mod, f, dtype)) |
| return result |
| elif isinstance(o, tvm.relay.backend.interpreter.ConstructorValue): |
| if o.constructor.name_hint == "Cons": |
| tl = vmobj_to_list(mod, o.fields[1], dtype) |
| hd = vmobj_to_list(mod, o.fields[0], dtype) |
| hd.extend(tl) |
| return hd |
| elif o.constructor.name_hint == "Nil": |
| return [] |
| elif "tensor_nil" in o.constructor.name_hint: |
| return [0] |
| elif "tensor" in o.constructor.name_hint: |
| return [o.fields[0].numpy()] |
| else: |
| raise RuntimeError("Unknown object type: %s" % o.constructor.name_hint) |
| else: |
| raise RuntimeError("Unknown object type: %s" % type(o)) |
| |
| |
| def check_tensor_array(ta_mod, ref_res, *args, dtype="float32", rtol=1e-5): |
| for kind in ["debug", "vm"]: |
| for target, dev in [("llvm", tvm.cpu(0))]: # testing.enabled_targets(): |
| if kind == "debug" and dev.device_type != tvm.cpu().device_type: |
| continue |
| result = relay.create_executor(kind, mod=ta_mod, device=dev, target=target).evaluate()( |
| *args |
| ) |
| got = vmobj_to_list(ta_mod, result, dtype) |
| tvm.testing.assert_allclose(ref_res, got, rtol=rtol, atol=rtol) |
| |
| |
| @tvm.testing.uses_gpu |
| def test_tensor_expand_dims(): |
| 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) |
| |
| run("float32") |
| run("int32") |
| |
| |
| @tvm.testing.uses_gpu |
| def test_tensor_array_constructor(): |
| 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) |
| |
| run("float32") |
| run("int32") |
| |
| |
| @tvm.testing.uses_gpu |
| def test_tensor_array_read(): |
| 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) |
| |
| run("float32") |
| run("int32") |
| |
| |
| @tvm.testing.uses_gpu |
| def test_tensor_array_write(): |
| 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) |
| |
| run("float32") |
| run("int32") |
| |
| |
| @tvm.testing.uses_gpu |
| def test_tensor_array_stack(): |
| 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) |
| |
| run("float32") |
| run("int32") |
| |
| |
| @tvm.testing.uses_gpu |
| def test_tensor_array_unstack(): |
| 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) |
| |
| run("float32") |
| run("int32") |
| |
| |
| @tvm.testing.uses_gpu |
| def test_tensor_take(): |
| 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) |
| |
| run("float32") |
| run("int32") |
| |
| |
| @tvm.testing.uses_gpu |
| def test_tensor_concatenate(): |
| 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) |
| |
| run("float32") |
| run("int32") |
| |
| |
| @tvm.testing.uses_gpu |
| def test_tensor_array_concat(): |
| 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) |
| |
| run("float32") |
| run("int32") |
| |
| |
| @tvm.testing.uses_gpu |
| def test_tensor_array_scatter(): |
| 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, |
| ) |
| |
| run("float32") |
| run("int32") |
| |
| |
| @tvm.testing.uses_gpu |
| def test_tensor_array_split(): |
| 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 |
| ) |
| |
| run("float32") |
| run("int32") |
| |
| |
| @tvm.testing.uses_gpu |
| def test_static_tensor_take(): |
| 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) |
| |
| run("float32", [10, 10]) |
| run("int32", [15, 11]) |
| |
| |
| @tvm.testing.uses_gpu |
| def test_static_tensor_concatenate(): |
| 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) |
| |
| run( |
| "float32", |
| [ |
| 5, |
| ], |
| ) |
| run("int32", [2, 3]) |
| |
| |
| @tvm.testing.uses_gpu |
| def test_static_tensor_expand_dims(): |
| 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) |
| |
| run("float32", []) |
| run( |
| "int32", |
| [ |
| 2, |
| ], |
| ) |
| |
| |
| @tvm.testing.uses_gpu |
| def test_static_tensor_array_constructor(): |
| def run(dtype, shape): |
| mod = tvm.IRModule() |
| p = Prelude(mod) |
| static_tensor_array_ops = StaticTensorArrayOps(p, dtype, shape) |
| static_tensor_array_ops.register() |
| tensor_constructor = p.get_name_static("tensor_constructor", dtype, shape) |
| assert tensor_constructor != None |
| |
| run("float32", [1, 1]) |
| |
| |
| @tvm.testing.uses_gpu |
| def test_static_tensor_array_read(): |
| 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) |
| 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], 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) |
| |
| run("float32", []) |
| run("int32", [2, 3]) |
| |
| |
| @tvm.testing.uses_gpu |
| def test_static_tensor_array_write(): |
| 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) |
| |
| run("float32", []) |
| run("int32", [2, 3]) |
| |
| |
| @tvm.testing.uses_gpu |
| def test_static_tensor_array_unstack(): |
| 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) |
| |
| run("float32", [4]) |
| run("int32", [2, 3]) |
| |
| |
| @tvm.testing.uses_gpu |
| def test_static_tensor_array_scatter(): |
| 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, |
| ) |
| |
| run("float32", [2, 3]) |
| run("int32", [2, 3]) |
| run( |
| "float32", |
| [2, 3], |
| [ |
| 2, |
| ], |
| ) |
| |
| |
| @tvm.testing.uses_gpu |
| def test_static_tensor_array_split(): |
| def run(dtype, shape, value_shape=None, lengths_shape=None): |
| mod = tvm.IRModule() |
| p = Prelude(mod) |
| static_tensor_array_ops = StaticTensorArrayOps(p, dtype, shape) |
| static_tensor_array_ops.register() |
| if value_shape is not None or lengths_shape is not None: |
| static_tensor_array_ops.define_tensor_array_split(value_shape, lengths_shape, False) |
| |
| # tensor array |
| v1 = relay.var("v1") |
| v2 = relay.var("v2") |
| v3 = relay.var("v2") |
| |
| adt_shape = [ |
| relay.Any(), |
| ] + shape[1:] |
| test_ops = StaticTensorArrayOps(p, dtype, adt_shape) |
| test_ops.register() |
| tensor_array = test_ops.get_global_var("tensor_array") |
| |
| tensor_array1 = tensor_array(relay.const(3)) |
| write_func = test_ops.get_global_var("tensor_array_write") |
| split_ops = StaticTensorArrayOps(p, dtype, shape) |
| split_ops.register() |
| split_func = split_ops.get_global_var("tensor_array_split") |
| tensor = p.get_tensor_ctor_static("tensor_constructor", dtype, test_ops.shape) |
| tensor_array1 = write_func(tensor_array1, 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)) |
| |
| # value tensor |
| value = relay.var("value") |
| |
| # lengths tensor |
| ta_len = relay.var("length") |
| |
| # create the split function |
| if value_shape is None: |
| tensor1 = p.get_tensor_ctor_static("tensor_constructor", dtype, shape) |
| else: |
| static_tensor_array_ops = StaticTensorArrayOps(p, dtype, value_shape) |
| static_tensor_array_ops.register() |
| tensor1 = p.get_tensor_ctor_static("tensor_constructor", dtype, test_ops.shape) |
| |
| tensor_array_split = split_func(tensor_array1, tensor1(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=value_shape or shape).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 |
| ) |
| |
| run("float32", [4, 3]) |
| run("int32", [4, 3]) |
| run( |
| "int32", |
| [relay.Any(), 3], |
| [4, 3], |
| [ |
| 2, |
| ], |
| ) |
| |
| |
| @tvm.testing.uses_gpu |
| def test_static_tensor_array_concat(): |
| 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) |
| |
| run("float32", [relay.Any(), 3]) |
| run("int32", [relay.Any(), 3]) |
| |
| |
| @tvm.testing.uses_gpu |
| def test_static_tensor_array_gather(): |
| 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) |
| |
| run("float32", []) |
| run("int32", [2, 3]) |
| |
| |
| @tvm.testing.uses_gpu |
| def test_static_tensor_array_stack(): |
| 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) |
| |
| run("float32", []) |
| run("int32", [2, 3]) |
| |
| |
| @tvm.testing.uses_gpu |
| def test_static_tensor_get_data(): |
| 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) |
| |
| run("float32", []) |
| run("int32", [2, 3]) |
| |
| |
| if __name__ == "__main__": |
| tvm.testing.main() |
