From 14d6c1736e5df12b47edbc3a2f59eecb0c696edb Mon Sep 17 00:00:00 2001
From: Chao Liu <monklof@gmail.com>
Date: Tue, 9 Mar 2021 04:37:36 +0800
Subject: [PATCH] [Relay] add ShapeFunc for one_hot op (#7490)

* [Relay] add ShapeFunc for one_hot op

* fix pylint

* add test for shapefunc of one_hot op
---
 3rdparty/vta-hw                   |  2 +-
 python/tvm/relay/op/_transform.py | 25 +++++++++++++++++++++++++
 tests/python/relay/test_any.py    | 22 ++++++++++++++++++++++
 3 files changed, 48 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/python/tvm/relay/op/_transform.py b/python/tvm/relay/op/_transform.py
index ad43192519d5..2bffd978d8b8 100644
--- a/python/tvm/relay/op/_transform.py
+++ b/python/tvm/relay/op/_transform.py
@@ -256,6 +256,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 b75cc5f5e750..32292de4c8ea 100644
--- a/tests/python/relay/test_any.py
+++ b/tests/python/relay/test_any.py
@@ -260,6 +260,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)