Op ReverseV2 (#594)

* RV2 -> RS for axes=[0, 1] and const shapes

* Working ReverseV2 for axis=0

* ReverseV2 op for constant axis

* ReverseV2 op for vector axis

* Support for empty axis for ReverseV2

* ReverseV2 op test with random tensor

* Minor fixes for ReverseV2

* ReverseV2 support for negative axis

* Computing the permutation for shape reset

* Working ReverseV2 with minimal Transpose nodes

* ReverseV2 test for 1D tensor with empty axis

* Removed unused variables

* changes after PR review v1

* minor changes fo ReverseV2 op support

* pylint fixes

* sort the axis vector

Author: satyajithj

tests/test_backend.py

@@ -2015,6 +2015,65 @@ def test_reverse_sequence_time_major(self):
         _ = tf.identity(x_, name=_TFOUTPUT)
         self._run_test_case([_OUTPUT], {_INPUT: x_val})
 
+
+    @check_opset_min_version(10, "ReverseSequence")
+    def test_reversev2_constant_axis(self):
+        # Tests for constant axis.
+        x_val_shape = [1, 2, 3, 4]
+        x_val = np.random.randint(0, 100, x_val_shape).astype(np.float32)
+        x = tf.placeholder(tf.float32, x_val_shape, name=_TFINPUT)
+        x_ = tf.reverse_v2(x, axis=[3])
+        _ = tf.identity(x_, name=_TFOUTPUT)
+        self._run_test_case([_OUTPUT], {_INPUT: x_val})
+        tf.reset_default_graph()
+
+        # Empty axis vector.
+        x_val_shape = [2, 3, 4]
+        x_val = np.random.randint(0, 100, x_val_shape).astype(np.float32)
+        x = tf.placeholder(tf.float32, x_val_shape, name=_TFINPUT)
+        x_ = tf.reverse_v2(x, axis=[])
+        _ = tf.identity(x_, name=_TFOUTPUT)
+        self._run_test_case([_OUTPUT], {_INPUT: x_val})
+
+
+    @check_opset_min_version(10, "ReverseSequence")
+    def test_reversev2_vector_axis(self):
+        x_val_shape = [1, 2, 3, 4]
+        x_val = np.random.randint(0, 100, x_val_shape).astype(np.float32)
+        x = tf.placeholder(tf.float32, x_val_shape, name=_TFINPUT)
+        x_ = tf.reverse_v2(x, axis=[0, -3, 2, 3])
+        _ = tf.identity(x_, name=_TFOUTPUT)
+        self._run_test_case([_OUTPUT], {_INPUT: x_val})
+        tf.reset_default_graph()
+
+        x_val_shape = [2, 3, 4]
+        x_val = np.random.randint(0, 100, x_val_shape).astype(np.float32)
+        x = tf.placeholder(tf.float32, x_val_shape, name=_TFINPUT)
+        x_ = tf.reverse_v2(x, axis=[-3, 1, 2])
+        _ = tf.identity(x_, name=_TFOUTPUT)
+        self._run_test_case([_OUTPUT], {_INPUT: x_val})
+        tf.reset_default_graph()
+
+        x_val_shape = [5, 5, 9, 7, 8, 9]
+        x_val = np.random.randint(0, 100, x_val_shape).astype(np.float32)
+        x = tf.placeholder(tf.float32, [5, 5, 9, 7, 8, 9], name=_TFINPUT)
+        x_ = tf.reverse_v2(x, axis=[0, 1, -2, 3, 5])
+        _ = tf.identity(x_, name=_TFOUTPUT)
+        self._run_test_case([_OUTPUT], {_INPUT: x_val})
+
+
+    @check_opset_min_version(10, "ReverseSequence")
+    def test_reversev2_1D_tensor(self):
+        # For tensors with 1 dimension and no axis to reverse.
+        # Adds an identity block.
+        x_val_shape = [4]
+        x_val = np.random.randint(0, 100, x_val_shape).astype(np.float32)
+        x = tf.placeholder(tf.float32, x_val_shape, name=_TFINPUT)
+        x_ = tf.reverse_v2(x, axis=[])
+        _ = tf.identity(x_, name=_TFOUTPUT)
+        self._run_test_case([_OUTPUT], {_INPUT: x_val})
+
+
     @check_opset_min_version(8, "where")
     def test_where(self):
         x_val = np.array([1, 2, -3, 4, -5, -6, -7, 8, 9, 0], dtype=np.float32)

tf2onnx/onnx_opset/tensor.py

@@ -1291,3 +1291,157 @@ def version_10(cls, ctx, node, **kwargs):
         target_dtype = TensorProto.INT64
         if seq_len_dtype != target_dtype:
             ctx.insert_new_node_on_input(node, "Cast", node.input[1], to=target_dtype)
+
+
+@tf_op("ReverseV2")
+class ReverseV2:
+    @classmethod
+    def version_10(cls, ctx, node, **kwargs):
+        # T output = ReverseV2(T input, int32|int64 seq_lengths, @int seq_dim, @int batch_dim)
+        # Implement tensorflow ReverseV2 op using multiple ReverseSequence (for each axis)
+        # and Transpose ops. We sort the axis vector (if non-empty) at the start. Each axis can
+        # be reversed only once (in tf) and so we can compute the transpose for each axis
+        # (other than 0), feed the tensor to a ReverseSequence node and finally transpose again
+        # to get back the original shape.
+
+        axes_node = node.inputs[1]
+        axes = axes_node.get_tensor_value(as_list=False)
+        # Current support is for when axis is a 1D tensor.
+        utils.make_sure(len(axes.shape) == 1 \
+            , "Currently no support for reverseV2 tensor axis")
+
+        axes = axes.tolist()
+        len_axes = len(axes)
+
+        # Store input and output parameters of the ReverseV2 node.
+        rv2_in_names = [node.input[0]]
+
+        input_shape = ctx.get_shape(node.input[0])
+        # Make sure input shape is not None
+        utils.make_sure(input_shape is not None, "shape of {} is None".format(node.input[0]))
+
+        input_rank = len(input_shape)
+
+        rv2_node_name = node.name
+        # ReverseV2 has a single output.
+        rv2_output_dtypes = node.output_dtypes
+        rv2_output_shapes = node.output_shapes
+
+        const_name_root = rv2_node_name + '_Const'
+
+        # Remove ReverseV2 node from graph.
+        ctx.remove_node(rv2_node_name)
+
+        # Variable to store input names for the next node.
+        inputs = rv2_in_names
+
+        new_node = None
+
+        # Empty axis vector.
+        if len_axes == 0:
+            # Replace ReverseV2 with an identity block.
+            new_node = ctx.make_node(
+                "Identity",
+                inputs=inputs,
+                outputs=node.output,
+                shapes=rv2_output_shapes,
+                dtypes=rv2_output_dtypes,
+                op_name_scope=rv2_node_name,
+            )
+
+        else:
+            # For negative indices use the positive counterpart.
+            for i, ax in enumerate(axes):
+                if ax < 0:
+                    axes[i] += input_rank
+
+            axes = sorted(axes)
+
+            orig_perm = list(range(input_rank))
+            curr_perm = []
+
+            # Add ReverseSequence nodes for each element of axis.
+            for i in range(len_axes):
+
+                axis = axes[i]
+
+                curr_perm = orig_perm.copy()
+                # Permutation indices relative to original tensor.
+                curr_perm[axis], curr_perm[0] = curr_perm[0], curr_perm[axis]
+
+                # Add a Transpose node if the axis != 0 (finish first due to sort).
+                if axis != 0:
+                    # Permutation indices for the transpose node relative to IN tensor shape.
+                    new_node = ctx.make_node(
+                        "Transpose",
+                        inputs=inputs,
+                        op_name_scope=rv2_node_name,
+                        dtypes=rv2_output_dtypes,
+                        attr={"perm": curr_perm}
+                    )
+
+                    inputs = [new_node.output[0]]
+
+                # Add a Constant node (seq_len) for ReverseSequence.
+
+                # Index 1 for the shape should not return 0
+                # since the input must have rank >= 2.
+                rs_batch_size = ctx.get_shape(inputs[-1])[1]
+
+                # Make sure rs_batch_size and input_shape[axis] are not -1 each
+                utils.make_sure(input_shape[axis] is not -1 \
+                    , "shape of axis {} is unknown".format(axis))
+                utils.make_sure(rs_batch_size is not -1 \
+                    , "ReverseSequence batch size for axis {} is unknown".format(axis))
+
+                seq_list = [input_shape[axis]] * rs_batch_size
+                seq_array = np.asarray(seq_list, dtype=np.int64) # dtype should be int64
+
+                const_seq_name = utils.make_name(const_name_root)
+                new_node = ctx.make_const(name=const_seq_name, np_val=seq_array)
+                inputs.append(new_node.output[0])
+
+                # Add a ReverseSequence node.
+
+                # If processing for the final axis and the tensor shape permutation is
+                # original then the output is fed to the output of the ReverseV2 node.
+                #
+                # Else a new output is created which is fed to a Transpose node.
+                rs_out_name = node.output if \
+                    ((i == len_axes - 1) and (curr_perm == orig_perm)) \
+                        else None
+
+                rs_out_shapes = None if rs_out_name is None else rv2_output_shapes
+
+                new_node = ctx.make_node(
+                    "ReverseSequence",
+                    inputs=inputs,
+                    op_name_scope=rv2_node_name,
+                    outputs=rs_out_name,
+                    shapes=rs_out_shapes,
+                    dtypes=rv2_output_dtypes,
+                    attr={"batch_axis": 1, "time_axis": 0}
+                )
+
+                inputs = [new_node.output[0]]
+
+            # Additional transpose block is required if the current
+            # permutation list is not the original one.
+            if curr_perm != orig_perm:
+
+                # Compute the required permutation list.
+                if len_axes != 1:
+                    for i, ax in enumerate(axes[::-1][1:]):
+                        curr_perm[0], curr_perm[ax] = \
+                            curr_perm[ax], curr_perm[0]
+
+                # Add a Transpose node to restore shape.
+                new_node = ctx.make_node(
+                    "Transpose",
+                    inputs=inputs,
+                    op_name_scope=rv2_node_name,
+                    outputs=node.output,
+                    shapes=rv2_output_shapes,
+                    dtypes=rv2_output_dtypes,
+                    attr={"perm": curr_perm}
+                )