From d46c5380bc5fb164a990f20347765069e0267653 Mon Sep 17 00:00:00 2001 From: monklof Date: Mon, 22 Feb 2021 15:52:38 +0800 Subject: [PATCH 1/3] [Relay] add ShapeFunc for one_hot op --- python/tvm/relay/op/_transform.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/python/tvm/relay/op/_transform.py b/python/tvm/relay/op/_transform.py index 01bcf4a6cf60..d24aef2028a2 100644 --- a/python/tvm/relay/op/_transform.py +++ b/python/tvm/relay/op/_transform.py @@ -220,6 +220,28 @@ def strided_slice_shape_func(attrs, inputs, _): ) ] +@script +def _one_hot_shape_func(indices_shape, depth, axis): + in_ndim = indices_shape.shape[0] + out_ndim = in_ndim + 1 + true_axis = in_ndim if axis == -1 else axis + indices_i = 0 + out = output_tensor((out_ndim, ), "int64") + for i in range(out_ndim): + if i == true_axis: + out[i] = int64(depth) + else: + out[i] = int64(indices_shape[indices_i]) + indices_i += 1 + return out + +@_reg.register_shape_func("one_hot", False) +def one_hot_shape_func(attrs, inputs, _): + """ + Shape func for one_hot + """ + shape_func = [_one_hot_shape_func(inputs[0],convert(attrs.depth), convert(attrs.axis))] + return shape_func @script def _concatenate_shape_func(inputs, axis): From 1b377a47b6a0b3c0c17167eb3d08fceeee844635 Mon Sep 17 00:00:00 2001 From: monklof Date: Tue, 23 Feb 2021 11:34:38 +0800 Subject: [PATCH 2/3] fix pylint --- python/tvm/relay/op/_transform.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/python/tvm/relay/op/_transform.py b/python/tvm/relay/op/_transform.py index d24aef2028a2..4c1f02a29aee 100644 --- a/python/tvm/relay/op/_transform.py +++ b/python/tvm/relay/op/_transform.py @@ -220,13 +220,14 @@ def strided_slice_shape_func(attrs, inputs, _): ) ] + @script def _one_hot_shape_func(indices_shape, depth, axis): in_ndim = indices_shape.shape[0] out_ndim = in_ndim + 1 true_axis = in_ndim if axis == -1 else axis indices_i = 0 - out = output_tensor((out_ndim, ), "int64") + out = output_tensor((out_ndim,), "int64") for i in range(out_ndim): if i == true_axis: out[i] = int64(depth) @@ -235,14 +236,16 @@ def _one_hot_shape_func(indices_shape, depth, axis): indices_i += 1 return out + @_reg.register_shape_func("one_hot", False) def one_hot_shape_func(attrs, inputs, _): """ Shape func for one_hot """ - shape_func = [_one_hot_shape_func(inputs[0],convert(attrs.depth), convert(attrs.axis))] + shape_func = [_one_hot_shape_func(inputs[0], convert(attrs.depth), convert(attrs.axis))] return shape_func + @script def _concatenate_shape_func(inputs, axis): ndim = inputs[0].shape[0] From 5216e14c3101c72e31b7e9482826426005a811af Mon Sep 17 00:00:00 2001 From: monklof Date: Mon, 8 Mar 2021 15:50:13 +0800 Subject: [PATCH 3/3] add test for shapefunc of one_hot op --- 3rdparty/vta-hw | 2 +- tests/python/relay/test_any.py | 22 ++++++++++++++++++++++ 2 files changed, 23 insertions(+), 1 deletion(-) diff --git a/3rdparty/vta-hw b/3rdparty/vta-hw index 57db5a718c74..87ce9acfae55 160000 --- a/3rdparty/vta-hw +++ b/3rdparty/vta-hw @@ -1 +1 @@ -Subproject commit 57db5a718c74a788c98120ebbe1230797be698c8 +Subproject commit 87ce9acfae550d1a487746e9d06c2e250076e54c diff --git a/tests/python/relay/test_any.py b/tests/python/relay/test_any.py index 9d05631a753a..4532ea301e31 100644 --- a/tests/python/relay/test_any.py +++ b/tests/python/relay/test_any.py @@ -239,6 +239,28 @@ def test_any_reshape(): verify_any_reshape(any_dims(3), (-4, 2, -1, -2), (6, 3, 4), (2, 3, 3, 4)) +def verify_any_one_hot(indices_shape, indices_np_shape, depth, on_value, off_value, axis, dtype): + indices = relay.var("indices", shape=indices_shape, dtype="int32") + on_value_const = relay.const(on_value, dtype) + off_value_const = relay.const(off_value, dtype) + y = relay.one_hot(indices, on_value_const, off_value_const, depth, axis=axis, dtype=dtype) + params = [indices] + mod = tvm.IRModule() + mod["main"] = relay.Function(params, y) + + indices_npy = np.random.randint(0, depth, size=indices_np_shape).astype("int32") + out_npy = tvm.topi.testing.one_hot(indices_npy, on_value, off_value, depth, axis, dtype) + args = [indices_npy] + check_result(args, mod, out_npy) + + +@tvm.testing.uses_gpu +def test_any_one_hot(): + verify_any_one_hot(any_dims(1), (3,), 3, 1, 0, -1, "int32") + verify_any_one_hot(any_dims(2), (2, 2), 5, 0.5, -0.5, 1, "float32") + verify_any_one_hot(any_dims(4), (3, 2, 4, 5), 6, 1.0, 0.0, 0, "float32") + + def verify_any_argwhere(x_shape, x_np_shape, dtype="bool"): x = relay.var("x", shape=x_shape, dtype=dtype) y = relay.argwhere(x)