Fix Roll op for negative shifts (#1616)

Signed-off-by: Tom Wildenhain <[email protected]>

Co-authored-by: Guenther Schmuelling <[email protected]>

Author: TomWildenhain-Microsoft

tests/test_backend.py

@@ -1084,6 +1084,16 @@ def func(input_ids):
         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_opset_min_version(10, "Slice")
+    def test_roll_neg_shift(self):
+        x_val = np.arange(4 * 3 * 5 * 2, dtype=np.float32).reshape((4, 3, 5, 2))
+        shift_val = np.array([-2, 13, -3], dtype=np.int32)
+        axes_val = np.array([1, 2, -1], dtype=np.int32)
+        def func(x, shift):
+            x_ = tf.roll(x, shift, axes_val)
+            return tf.identity(x_, name=_TFOUTPUT)
+        self._run_test_case(func, [_OUTPUT], {_INPUT: x_val, _INPUT1: shift_val})
+
     @check_tf_min_version("2.2")
     def test_large_model_format(self):
         x_val = np.array([2.0], dtype=np.float32)

tf2onnx/onnx_opset/tensor.py

@@ -369,10 +369,11 @@ 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]
+        rank = ctx.get_rank(node.input[0])
+        axes = [a if a >= 0 else a + rank for a in axes]
+
         shifts_dtype = ctx.get_dtype(node.input[1])
         if shifts_dtype != TensorProto.INT64:
             shifts_casted = ctx.insert_new_node_on_input(node, "Cast", node.input[1], to=TensorProto.INT64).output[0]
@@ -395,7 +396,8 @@ def any_version(cls, opset, ctx, node, **kwargs):
         for axis, shift in zip(axes, shifts_split):
             len_along_axis = GraphBuilder(ctx).make_slice(
                 {"data": shape_node.output[0], "ends": [axis + 1], "starts": [axis]})
-            remaining_len = ctx.make_node("Sub", [len_along_axis, shift], op_name_scope=node.name).output[0]
+            shift_mod = ctx.make_node("Mod", [shift, len_along_axis]).output[0]
+            remaining_len = ctx.make_node("Sub", [len_along_axis, shift_mod], op_name_scope=node.name).output[0]
             axes_const = ctx.make_const(utils.make_name("axes_const"), np.array([axis], np.int64)).output[0]
             slice_one = ctx.make_node("Slice", [data, zero_const, remaining_len, axes_const], op_name_scope=node.name)
             slice_two = ctx.make_node("Slice", [data, remaining_len, len_along_axis, axes_const],