Merge pull request #589 from lucienwang1009/ellipse_mask

fix ellipsis_mask bug

Author: wayuanho

tests/test_backend.py

@@ -1621,6 +1621,28 @@ def test_strided_slice7(self):
         _ = tf.identity(x_, name=_TFOUTPUT)
         self._run_test_case([_OUTPUT], {_INPUT: x_val})
 
+    @skip_caffe2_backend("multiple dims not supported")
+    def test_strided_slice8(self):
+        x_val = np.arange(1 * 2 * 3 * 4 * 5 * 6).astype("float32").reshape((1, 2, 3, 4, 5, 6))
+        x = tf.placeholder(tf.float32, x_val.shape, name=_TFINPUT)
+        x_ = x[0:1, ..., 1, 2:, :6]
+        _ = tf.identity(x_, name=_TFOUTPUT)
+        self._run_test_case([_OUTPUT], {_INPUT: x_val})
+
+        tf.reset_default_graph()
+        x_val = np.arange(1 * 2 * 3 * 4 * 5 * 6).astype("float32").reshape((1, 2, 3, 4, 5, 6))
+        x = tf.placeholder(tf.float32, x_val.shape, name=_TFINPUT)
+        x_ = x[0:1, 1, 2:, :6, ...]
+        _ = tf.identity(x_, name=_TFOUTPUT)
+        self._run_test_case([_OUTPUT], {_INPUT: x_val})
+
+        tf.reset_default_graph()
+        x_val = np.arange(1 * 2 * 3 * 4 * 5 * 6).astype("float32").reshape((1, 2, 3, 4, 5, 6))
+        x = tf.placeholder(tf.float32, x_val.shape, name=_TFINPUT)
+        x_ = x[..., 0:1, 1, 2:, :6]
+        _ = tf.identity(x_, name=_TFOUTPUT)
+        self._run_test_case([_OUTPUT], {_INPUT: x_val})
+
     @check_opset_min_version(10, "Slice")
     @skip_caffe2_backend("multiple dims not supported")
     def test_strided_slice_dynamic_1(self):
@@ -1702,6 +1724,35 @@ def test_strided_slice_dynamic_6(self):
         _ = tf.identity(x_, name=_TFOUTPUT)
         self._run_test_case([_OUTPUT], {_INPUT: x_val, _INPUT1: y_val})
 
+    @check_opset_min_version(10, "Slice")
+    @skip_caffe2_backend("multiple dims not supported")
+    def test_strided_slice_dynamic_7(self):
+        x_val = np.arange(1 * 2 * 3 * 4 * 5 * 6).astype("float32").reshape((1, 2, 3, 4, 5, 6))
+        y_val = np.array(1, dtype=np.int32)
+        x = tf.placeholder(tf.float32, x_val.shape, name=_TFINPUT)
+        y = tf.placeholder(tf.int32, y_val.shape, name=_TFINPUT1)
+        x_ = x[0:y, ..., y, y:, :y]
+        _ = tf.identity(x_, name=_TFOUTPUT)
+        self._run_test_case([_OUTPUT], {_INPUT: x_val, _INPUT1: y_val})
+
+        tf.reset_default_graph()
+        x_val = np.arange(1 * 2 * 3 * 4 * 5 * 6).astype("float32").reshape((1, 2, 3, 4, 5, 6))
+        y_val = np.array(1, dtype=np.int32)
+        x = tf.placeholder(tf.float32, x_val.shape, name=_TFINPUT)
+        y = tf.placeholder(tf.int32, y_val.shape, name=_TFINPUT1)
+        x_ = x[0:y, y, y:, :y, ...]
+        _ = tf.identity(x_, name=_TFOUTPUT)
+        self._run_test_case([_OUTPUT], {_INPUT: x_val, _INPUT1: y_val})
+
+        tf.reset_default_graph()
+        x_val = np.arange(1 * 2 * 3 * 4 * 5 * 6).astype("float32").reshape((1, 2, 3, 4, 5, 6))
+        y_val = np.array(1, dtype=np.int32)
+        x = tf.placeholder(tf.float32, x_val.shape, name=_TFINPUT)
+        y = tf.placeholder(tf.int32, y_val.shape, name=_TFINPUT1)
+        x_ = x[..., 0:y, y, y:, :y]
+        _ = tf.identity(x_, name=_TFOUTPUT)
+        self._run_test_case([_OUTPUT], {_INPUT: x_val, _INPUT1: y_val})
+
     @skip_caffe2_backend("fails with schema error")
     @check_opset_min_version(7, "batchnorm")
     def test_batchnorm(self):

tf2onnx/onnx_opset/tensor.py

@@ -567,6 +567,7 @@ def version_1(cls, ctx, node, **kwargs):
             attr = node.get_attr(attr_name)
             if attr is not None and attr.i != 0:
                 raise ValueError("StridedSlice: attribute " + attr_name + " not supported")
+
         onnx_dtype = ctx.get_dtype(node.input[1])
         np_dtype = utils.ONNX_TO_NUMPY_DTYPE[onnx_dtype]
         max_size = np.iinfo(np_dtype).max
@@ -586,17 +587,27 @@ def version_1(cls, ctx, node, **kwargs):
         axes = []
         # onnx slice op can't remove a axis, track axis and add a squeeze op if needed
         needs_squeeze = []
+        # ellipsis: one bit at most can be 1. An ellipsis implicitly creates as many range specifications as
+        # necessary to fully specify the sliced range for every dimension.
+        # For example for a 4-dimensional tensor foo the slice foo[2, ..., 5:8] implies foo[2, :, :, 5:8]
+        # NOTE: we ignore those axes denoted by ellipsis using `axes` attribute
+        ellipsis_gap = 0
         for idx, begin_item in enumerate(begin):
-            end_item = end[idx]
             if strides[idx] != 1:
                 raise ValueError("StridedSlice: only strides=1 is supported")
-            axes.append(idx)
-
             if (ellipsis_mask >> idx) & 1:
-                new_begin.append(0)
-                new_end.append(max_size)
+                input_shape = ctx.get_shape(node.input[0])
+                utils.make_sure(
+                    input_shape is not None,
+                    "StridedSlice op {} requires the shape of input".format(node.name)
+                )
+                ellipsis_gap = len(input_shape) - len(begin)
                 continue
 
+            # ignore ellipsis axes
+            axes.append(idx + ellipsis_gap)
+            end_item = end[idx]
+
             # an implicit condition is stride == 1 (checked in above)
             if begin_item < 0 and end_item == 0:
                 end_item = max_size
@@ -606,7 +617,7 @@ def version_1(cls, ctx, node, **kwargs):
                 new_begin.append(begin_item)
                 end_item = begin_item + 1 if begin_item != -1 else max_size
                 new_end.append(end_item)
-                needs_squeeze.append(idx)
+                needs_squeeze.append(idx + ellipsis_gap)
                 continue
 
             mask = (begin_mask >> idx) & 1
@@ -698,30 +709,46 @@ def version_10(cls, ctx, node, **kwargs):
         ellipsis_mask = ellipsis_mask.i if ellipsis_mask is not None else 0
         shrink_axis_mask = node.get_attr("shrink_axis_mask")
         shrink_axis_mask = shrink_axis_mask.i if shrink_axis_mask is not None else 0
-        input_shape = ctx.get_shape(node.input[1]) or \
+        param_shape = ctx.get_shape(node.input[1]) or \
             ctx.get_shape(node.input[2]) or \
             ctx.get_shape(node.input[3])
-        utils.make_sure(input_shape, "StridedSlice op {} requires the shape of begin/end/strides".format(node.name))
-        input_rank = input_shape[0]
+        utils.make_sure(
+            param_shape is not None,
+            "StridedSlice op {} requires the shape of begin/end/strides".format(node.name)
+        )
+        param_rank = param_shape[0]
         # use in onnx graph to mask begin
-        new_begin_mask = [1] * input_rank
+        new_begin_mask = [1] * param_rank
         # use in onnx graph to mask end
-        new_end_mask = [min_size] * input_rank
+        new_end_mask = [min_size] * param_rank
         # for shrink mask, if shrink mask is 1, set stride to be max_size
-        shrink_strided_mask = [min_size] * input_rank
+        shrink_strided_mask = [min_size] * param_rank
+        axes = []
         # onnx slice op can't remove a axis, track axis and add a squeeze op if needed
         needs_squeeze = []
-        for idx in range(input_rank):
+        ellipsis_gap = 0
+        for idx in range(param_rank):
             if (ellipsis_mask >> idx) & 1:
+                input_shape = ctx.get_shape(node.input[0])
+                utils.make_sure(
+                    input_shape is not None,
+                    "StridedSlice op {} requires the shape of input".format(node.name)
+                )
+                ellipsis_gap = len(input_shape) - param_rank
+                # handle the redundant param
                 new_begin_mask[idx] = 0
                 new_end_mask[idx] = max_size
+                axes.append(idx)
                 continue
 
+            # ignore ellipsis axes
+            axes.append(idx + ellipsis_gap)
+
             mask = (shrink_axis_mask >> idx) & 1
             if mask != 0:
                 shrink_strided_mask[idx] = max_size
                 new_end_mask[idx] = max_size
-                needs_squeeze.append(idx)
+                needs_squeeze.append(idx + ellipsis_gap)
                 continue
 
             mask = (begin_mask >> idx) & 1
@@ -794,7 +821,7 @@ def version_10(cls, ctx, node, **kwargs):
         # create axes input
         axes_const = ctx.make_const(
             utils.make_name("slice_axes"),
-            np.arange(input_rank, dtype=np_dtype)
+            np.array(axes, dtype=np_dtype)
         )
         axes_output = axes_const.output[0]