Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[TFLite] Implemented REVERSE_V2 Operator for TFLite. #6304

Merged
merged 1 commit into from
Aug 21, 2020
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 17 additions & 0 deletions python/tvm/relay/frontend/tflite.py
Original file line number Diff line number Diff line change
Expand Up @@ -136,6 +136,7 @@ def __init__(self, model, subgraph, exp_tab):
'ROUND': self.convert_round,
'RSQRT': self.convert_rsqrt,
'REVERSE_SEQUENCE': self.convert_reverse_sequence,
'REVERSE_V2': self.convert_reverse_v2,
'SELECT': self.convert_select,
'SHAPE': self.convert_shape,
'SIN': self.convert_sin,
Expand Down Expand Up @@ -2973,6 +2974,22 @@ def convert_one_hot(self, op):

return out

def convert_reverse_v2(self, op):
"""Convert TFLite REVERSE_V2"""
input_tensors = self.get_input_tensors(op)
assert len(input_tensors) == 2, "input tensor's length should be 2"

input_expr = self.get_expr(input_tensors[0].tensor_idx)

# Getting axis value
axis = self.get_tensor_value(input_tensors[1])
if isinstance(axis, np.ndarray):
assert len(axis) == 1, "TFLite does not support multi-axis yet"
axis = int(axis)

out = _op.reverse(input_expr, axis)
return out


def get_expr(self, input_tensor_idx):
return self.exp_tab.get_expr(get_tensor_name(self.subgraph, input_tensor_idx))
Expand Down
27 changes: 27 additions & 0 deletions tests/python/frontend/tflite/test_forward.py
Original file line number Diff line number Diff line change
Expand Up @@ -2620,6 +2620,32 @@ def test_forward_fully_connected():
_test_fully_connected([5, 1, 1, 150], [150, 100], [100])


#######################################################################
# REVERSE_V2
# ----------

def _test_reverse_v2(input_shape, axis, dtype):
""" One iteration of REVERSE_V2 """
with tf.Graph().as_default():
input = np.random.randint(0, 100, size=input_shape).astype(dtype)
in_input = tf.placeholder(dtype=input.dtype, shape=input.shape, name="input")
in_axis = ops.convert_to_tensor(axis, dtype=axis.dtype)

out = array_ops.reverse(in_input, in_axis)

compare_tflite_with_tvm(
[input],
["input"],
[in_input],
[out])

def test_forward_reverse_v2():
""" REVERSE_V2 """
for dtype in ['float32', 'int32']:
_test_reverse_v2((5), np.array([0], dtype='int32'), dtype)
_test_reverse_v2((5, 6, 4, 2), np.array([2], dtype='int32'), dtype)


#######################################################################
# Custom Operators
# ----------------
Expand Down Expand Up @@ -3098,6 +3124,7 @@ def test_forward_mediapipe_hand_landmark():
test_forward_quantize_dequantize()
test_forward_arg_min_max()
test_forward_expand_dims()
test_forward_reverse_v2()

# NN
test_forward_convolution()
Expand Down