Merge pull request #388 from jiafatom/ellipse_mask

Support ellipse mask in stridedslice op

Author: guschmue

tests/test_backend.py

@@ -1243,6 +1243,12 @@ def test_strided_slice7(self):
         _ = tf.identity(x_, name=_TFOUTPUT)
         self._run_test_case([_OUTPUT], {_INPUT: x_val})
 
+        tf.reset_default_graph()
+        x = tf.placeholder(tf.float32, x_val.shape, name=_TFINPUT)
+        x_ = tf.strided_slice(x, [0, 1], [3, 4], [1, 1], ellipsis_mask=2)
+        _ = tf.identity(x_, name=_TFOUTPUT)
+        self._run_test_case([_OUTPUT], {_INPUT: x_val})
+
     @skip_caffe2_backend("fails with schema error")
     @check_opset_min_version(7, "batchnorm")
     def test_batchnorm(self):

tf2onnx/tfonnx.py

@@ -926,7 +926,7 @@ def expanddims_op7(ctx, node, name, args):
 
 def stridedslice_op(ctx, node, name, args):
     # for now we implement common cases. Things like strides!=1 are not mappable to onnx.
-    not_supported_attr = ["ellipsis_mask", "new_axis_mask"]
+    not_supported_attr = ["new_axis_mask"]
     for attr_name in not_supported_attr:
         attr = node.get_attr(attr_name)
         if attr is not None and attr.i != 0:
@@ -942,6 +942,8 @@ def stridedslice_op(ctx, node, name, args):
     begin_mask = begin_mask.i if begin_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
+    ellipsis_mask = node.get_attr("ellipsis_mask")
+    ellipsis_mask = ellipsis_mask.i if ellipsis_mask is not None else 0
     new_begin = []
     new_end = []
     axes = []
@@ -953,6 +955,11 @@ def stridedslice_op(ctx, node, name, args):
             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)
+            continue
+
         # an implicit condition is stride == 1 (checked in above)
         if begin_item < 0 and end_item == 0:
             end_item = max_size