Skip to content

Commit

Permalink
fix tf.roll with axes=-1 (#1338)
Browse files Browse the repository at this point in the history
Signed-off-by: Guenther Schmuelling <guschmue@microsoft.com>
  • Loading branch information
guschmue authored Feb 13, 2021
1 parent a4aeb1a commit c24dc6f
Show file tree
Hide file tree
Showing 2 changed files with 12 additions and 0 deletions.
10 changes: 10 additions & 0 deletions tests/test_backend.py
Original file line number Diff line number Diff line change
Expand Up @@ -911,6 +911,15 @@ def func(x, shift):
return tf.identity(x_, name=_TFOUTPUT)
self._run_test_case(func, [_OUTPUT], {_INPUT: x_val, _INPUT1: shift_val})

@check_opset_min_version(10, "Slice")
def test_roll_neg_axis(self):
def func(input_ids):
shifted_input_ids = tf.cast(input_ids, tf.int32)
shifted_input_ids = tf.roll(shifted_input_ids, 1, axis=-1)
return tf.identity(shifted_input_ids, name=_TFOUTPUT)
x_val = np.array([[0, 1, 2, 3, 4, 5, 6, 7], [1, 2, 3, 4, 5, 6, 7, 8]], dtype=np.int64)
self._run_test_case(func, [_OUTPUT], {_INPUT: x_val})

@check_tf_min_version("2.2")
def test_large_model_format(self):
x_val = np.array([2.0], dtype=np.float32)
Expand Down Expand Up @@ -4581,5 +4590,6 @@ def func(x):
self._run_test_case(func, [_OUTPUT], {_INPUT: x_val})



if __name__ == '__main__':
unittest_main()
2 changes: 2 additions & 0 deletions tf2onnx/onnx_opset/tensor.py
Original file line number Diff line number Diff line change
Expand Up @@ -356,6 +356,8 @@ class Roll:
def any_version(cls, opset, ctx, node, **kwargs):
utils.make_sure(node.inputs[2].is_const(), "Can only convert Roll is axis is const")
axes = node.inputs[2].get_tensor_value()
if axes == -1:
axes = len(ctx.get_shape(node.input[0])) + axes
if not isinstance(axes, list):
axes = [axes]
shifts_dtype = ctx.get_dtype(node.input[1])
Expand Down

0 comments on commit c24dc6f

Please sign in to comment.