[GitHub] [tvm] codeislife99 commented on a change in pull request #7441: [Frontend][Tensorflow] Add unique operator

[GitBox <gi...@apache.org>]

codeislife99 commented on a change in pull request #7441:
URL: https://github.com/apache/tvm/pull/7441#discussion_r574219984



##########
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:
       You can put this in the same shape func, You can return a tuple from the shape func and remove the list from this function. 




----------------------------------------------------------------
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:
users@infra.apache.org