ymwangg commented on a change in pull request #7441:
URL: https://github.com/apache/tvm/pull/7441#discussion_r574883580
##########
File path: python/tvm/relay/op/_transform.py
##########
@@ -885,3 +886,29 @@ def where_shape_func(attrs, inputs, _):
out_shape = _broadcast_shape_tensors(bcast_shape, cond_shape)
return [out_shape]
+
+
+register_strategy("unique", strategy.unique_strategy)
+register_pattern("unique", OpPattern.OPAQUE)
+
+
+@script
+def _unique_shape_1(data_shape):
+ shape_tensor = output_tensor((1,), "int64")
+ shape_tensor[0] = int64(data_shape[0])
+ return shape_tensor
+
+
+@script
+def _unique_shape_2(inputs):
+ shape_tensor = output_tensor((1,), "int64")
+ shape_tensor[0] = int64(1)
+ return shape_tensor
+
+
+@_reg.register_shape_func("unique", False)
+def unique_shape_func(attrs, inputs, _):
+ """
+ Shape func for unique operator.
+ """
+ return [_unique_shape_1(inputs[0]), _unique_shape_1(inputs[0]),
_unique_shape_1(inputs[0]), _unique_shape_2(inputs[0])]
Review comment:
Sure, I'll make it simpler.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
