Merge pull request #828 from jignparm/jignparm/fix_strided_slice_reverse

Fix strided slice, reverse stride with unbounded begin and end

Author: jignparm

tests/test_backend.py

@@ -1787,6 +1787,29 @@ def func(x, y):
             return tf.identity(x_, name=_TFOUTPUT)
         self._run_test_case(func, [_OUTPUT], {_INPUT: x_val, _INPUT1: y_val})
 
+    @check_opset_min_version(10, "Slice")
+    def test_strided_slice_reverse_1(self):
+        x_val = np.arange(16 * 32).astype(np.float32).reshape((1, 16, 32, 1))
+        def func(x):
+            return tf.concat([x[:, :, :10], x[:, :, :21:-1]], axis=0, name=_TFOUTPUT)
+        self._run_test_case(func, [_OUTPUT], {_INPUT: x_val})
+
+    @check_opset_min_version(10, "Slice")
+    def test_strided_slice_reverse_2(self):
+        x_val = np.arange(16 * 32).astype(np.float32).reshape((1, 16, 32, 1))
+        def func(x):
+            return tf.concat([x[:, :, :10], x[:, :, 9::-1]], axis=0, name=_TFOUTPUT)
+        self._run_test_case(func, [_OUTPUT], {_INPUT: x_val})
+
+    @check_opset_min_version(10, "Slice")
+    def test_strided_slice_reverse_3(self):
+        x_val = np.zeros((1, 16, 32, 1)).astype(np.float32)
+        y_val = np.array(9).astype(np.int32)
+        z_val = np.array(-1).astype(np.int32)
+        def func(x, y, z):
+            return tf.concat([x[:, :, :10], x[:, :, y::z]], axis=0, name=_TFOUTPUT)
+        self._run_test_case(func, [_OUTPUT], {_INPUT: x_val, _INPUT1: y_val, _INPUT2: z_val})
+
     @check_opset_min_version(10, "Slice")
     def test_new_axis_mask(self):
         def func(x, y):

tf2onnx/onnx_opset/math.py

@@ -97,8 +97,8 @@ def make_min_or_max_op(ctx, op_type, inputs, outputs,
         ctx.set_dtype(cast_node.output[0], origin_dtype)
         ctx.copy_shape(node.output[0], cast_node.output[0])
         actual_outputs = cast_node.output
-    ctx.make_node("Identity", actual_outputs, outputs=outputs,
-                  shapes=output_shapes, dtypes=output_dtypes)
+    final_node = ctx.make_node("Identity", actual_outputs, outputs=outputs,
+                               shapes=output_shapes, dtypes=output_dtypes)
 
     # tensorflow minimum/maximum does support broadcast, onnx < opset 8 does not.
     # handle this by doing something like:
@@ -124,6 +124,7 @@ def make_min_or_max_op(ctx, op_type, inputs, outputs,
             add_node = ctx.make_node("Add", [input_node.output[0], sub_node.output[0]],
                                      op_name_scope=input_node.name)
             node.input[i] = add_node.output[0]
+    return final_node
 
 
 @tf_op("Minimum", onnx_op="Min")

tf2onnx/onnx_opset/tensor.py

@@ -845,37 +845,51 @@ def version_10(cls, ctx, node, **kwargs):
         # mask begin
         new_begin_mask = np.array(new_begin_mask, dtype=np_dtype)
         if not np.all(new_begin_mask == 1):
-            if begin.is_const():
-                begin = ctx.make_const(
-                    utils.make_name("begin_masked"),
-                    begin.get_tensor_value(as_list=False) * new_begin_mask
-                )
+            if begin.is_const() and strides.is_const():
+                new_begin_vals = np.copy(begin.get_tensor_value(as_list=False))
+                strides_vals = strides.get_tensor_value(as_list=False)
+                idx1 = np.where(new_begin_mask == 0)
+                idx2 = np.where(strides_vals < 0)
+                idx3 = np.intersect1d(idx1, idx2)
+                new_begin_vals[idx3] = max_size
+                begin = ctx.make_const(utils.make_name("begin_masked"), new_begin_vals)
             else:
-                begin_mask_const = ctx.make_const(
-                    utils.make_name("begin_mask"),
-                    new_begin_mask
-                )
-                begin = ctx.make_node(
-                    "Mul", [begin.output[0], begin_mask_const.output[0]],
-                    op_name_scope=node.name
-                )
+                begin_mask_const = ctx.make_const(utils.make_name("begin_mask"), np.equal(new_begin_mask, 0))
+                zero_const = ctx.make_const(utils.make_name("zero_const"), np.zeros(1, dtype=np_dtype))
+                max_const = ctx.make_const(utils.make_name("max_const"), np.array(max_size, dtype=np_dtype))
+                op1 = ctx.make_node("Less", [strides.output[0], zero_const.output[0]], op_name_scope=node.name)
+                op2 = ctx.make_node("And", [op1.output[0], begin_mask_const.output[0]], op_name_scope=node.name)
+                begin = ctx.make_node("Where", [op2.output[0], max_const.output[0], begin.output[0]],
+                                      op_name_scope=node.name)
+
         # mask end
         new_end_mask = np.array(new_end_mask, dtype=np_dtype)
         end_output = end.output[0]
         if not np.all(new_end_mask == min_size):
-            if end.is_const():
-                end = ctx.make_const(
-                    utils.make_name("end_masked"),
-                    np.maximum(end.get_tensor_value(as_list=False), new_end_mask)
-                )
+            if end.is_const() and strides.is_const():
+                new_end_mask = np.maximum(end.get_tensor_value(as_list=False), new_end_mask)
+                idx = np.where(new_end_mask == max_size)
+                sign = np.sign(strides.get_tensor_value(as_list=False))[idx]
+                new_end_mask[idx] = new_end_mask[idx] * sign
+                end = ctx.make_const(utils.make_name("end_masked"), new_end_mask)
                 end_output = end.output[0]
             else:
-                end_mask_const = ctx.make_const(
-                    utils.make_name("end_mask"),
-                    np.array(new_end_mask, dtype=np_dtype)
-                )
-                end_output = utils.make_name("{}__end".format(node.name))
-                math.make_min_or_max_op(ctx, "Max", [end.output[0], end_mask_const.output[0]], [end_output])
+                # Overlay new_end_mask with specified end values.
+                # Adjust max_size to min_size if steps are < 0
+                max_const = ctx.make_const(utils.make_name("max_const"), np.array(max_size, dtype=np_dtype))
+                min_const = ctx.make_const(utils.make_name("min_const"), np.array(min_size, dtype=np_dtype))
+                zero_const = ctx.make_const(utils.make_name("zero_const"), np.zeros(1, dtype=np_dtype))
+                end_mask_const = ctx.make_const(utils.make_name("end_mask"), np.array(new_end_mask, dtype=np_dtype))
+                outputname = utils.make_name("{}__newendmask".format(node.name))
+                new_end_mask = math.make_min_or_max_op(ctx, "Max", [end.output[0], end_mask_const.output[0]],
+                                                       [outputname])
+                op1 = ctx.make_node("Less", [strides.output[0], zero_const.output[0]], op_name_scope=node.name)
+                op2 = ctx.make_node("Equal", [new_end_mask.output[0], max_const.output[0]], op_name_scope=node.name)
+                op3 = ctx.make_node("And", [op2.output[0], op1.output[0]], op_name_scope=node.name)
+                final_end = ctx.make_node("Where", [op3.output[0], min_const.output[0],
+                                                    new_end_mask.output[0]], op_name_scope=node.name)
+                end_output = final_end.output[0]
+
         # mask strides for shrink
         shrink_strided_mask = np.array(shrink_strided_mask, dtype=np_dtype)
         strides_output = strides.output[0]