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/python/tvm/relay/op/_transform.py b/python/tvm/relay/op/_transform.py index 01bcf4a6cf60..4c1f02a29aee 100644 --- a/python/tvm/relay/op/_transform.py +++ b/python/tvm/relay/op/_transform.py @@ -221,6 +221,31 @@ 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): ndim = inputs[0].shape[0] 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)