Merge pull request #774 from RandySheriffH/rashuai/xlnet

Rashuai/xlnet

Author: guschmue

tests/test_backend.py

@@ -1828,6 +1828,16 @@ def test_strided_slice_dynamic_7(self):
         _ = tf.identity(x_, name=_TFOUTPUT)
         self._run_test_case([_OUTPUT], {_INPUT: x_val, _INPUT1: y_val})
 
+    @check_opset_min_version(10, "Slice")
+    def test_new_axis_mask(self):
+        x_val = np.arange(5*10*10*10*10*20*30).astype("float32").reshape((5, 10, 10, 10, 10, 20, 30))
+        y_val = np.array(9, 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[tf.newaxis, 0:y, y::2, tf.newaxis, :, tf.newaxis, :y, tf.newaxis, ..., 9]
+        _ = 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

@@ -738,19 +738,19 @@ def version_10(cls, ctx, node, **kwargs):
         #                 @int shrink_axis_mask, @int new_axis_mask)
         # T output = Slice(T input, Tind starts, Tind ends, Tind axes, Tind steps)
         # "ends" are exclusive, "axes" and "steps" are optional, their default val are [0, ...] and 1
+        input_x = node.inputs[0]
         begin = node.inputs[1]
         end = node.inputs[2]
         strides = node.inputs[3]
+        new_axis_mask = node.get_attr("new_axis_mask")
+        new_axis_mask = new_axis_mask.i if new_axis_mask is not None else 0
+
         if begin.is_const() and end.is_const() and strides.is_const() \
-                and all(val == 1 for val in strides.get_tensor_value()):
+                and all(val == 1 for val in strides.get_tensor_value()) \
+                and new_axis_mask == 0:
             cls.version_1(ctx, node, **kwargs)
             return
 
-        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:
-                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]
 
@@ -769,6 +769,15 @@ 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
+        if new_axis_mask != 0:
+            unqueeze_at = []
+            for bit in range(32):
+                if (new_axis_mask >> bit) & 1 == 1:
+                    unqueeze_at.append(bit)
+                    begin_mask |= 1 << bit
+                    end_mask |= 1 << bit
+            input_x = ctx.make_node("Unsqueeze", [input_x.output[0]], {"axes": unqueeze_at})
+
         param_shape = ctx.get_shape(node.input[1]) or \
                       ctx.get_shape(node.input[2]) or \
                       ctx.get_shape(node.input[3])
@@ -789,7 +798,7 @@ def version_10(cls, ctx, node, **kwargs):
         ellipsis_gap = 0
         for idx in range(param_rank):
             if (ellipsis_mask >> idx) & 1:
-                input_shape = ctx.get_shape(node.input[0])
+                input_shape = ctx.get_shape(input_x.output[0])
                 utils.make_sure(
                     input_shape is not None,
                     "StridedSlice op {} requires the shape of input".format(node.name)
@@ -886,7 +895,7 @@ def version_10(cls, ctx, node, **kwargs):
         axes_output = axes_const.output[0]
 
         inputs_map = {
-            "data": node.input[0],
+            "data": input_x.output[0],
             "starts": begin.output[0],
             "ends": end_output,
             "steps": strides_output,