Merge branch 'master' into tom/RaggedGather

Author: TomWildenhain-Microsoft

tests/test_backend.py

@@ -871,6 +871,26 @@ def func(x):
             return tf.identity(x_, name=_TFOUTPUT)
         self._run_test_case(func, [_OUTPUT], {_INPUT: x_val})
 
+    @check_tf_min_version("1.14")
+    @check_opset_min_version(12, "GatherND with batch_dims")
+    def test_gather_batch_dims_no_trans(self):
+        x_val = np.arange(2 * 2 * 3 * 5 * 4, dtype=np.float32).reshape((2, 2, 3, 5, 4))
+        idx_val = np.array([[[1, 0, 2, 0], [1, 1, 1, 0]], [[0, 0, 0, 0], [2, 1, 1, 0]]], dtype=np.int32)
+        def func(x, idx):
+            x_ = tf.gather(x, idx, batch_dims=2, axis=2)
+            return tf.identity(x_, name=_TFOUTPUT)
+        self._run_test_case(func, [_OUTPUT], {_INPUT: x_val, _INPUT1: idx_val})
+
+    @check_tf_min_version("1.14")
+    @check_opset_min_version(12, "GatherND with batch_dims")
+    def test_gather_batch_dims(self):
+        x_val = np.arange(2 * 2 * 3 * 5 * 4, dtype=np.float32).reshape((2, 2, 3, 5, 4))
+        idx_val = np.array([[[1, 0, 2, 0], [1, 1, 1, 0]], [[0, 0, 0, 0], [2, 1, 1, 0]]], dtype=np.int32)
+        def func(x, idx):
+            x_ = tf.gather(x, idx, batch_dims=2, axis=3)
+            return tf.identity(x_, name=_TFOUTPUT)
+        self._run_test_case(func, [_OUTPUT], {_INPUT: x_val, _INPUT1: idx_val})
+
     @check_opset_min_version(10, "Slice")
     def test_roll_axis_scalar(self):
         x_val = np.arange(4 * 3 * 5 * 2, dtype=np.float32).reshape((4, 3, 5, 2))
@@ -3852,6 +3872,28 @@ def func(splits, rt_dense_values, indices):
         self._run_test_case(func, [_OUTPUT, _OUTPUT1],
                             {_INPUT: splits_val, _INPUT1: dense_vals_val, _INPUT2: indices_val})
 
+    def test_ragged_tensor_to_tensor(self):
+        splits_val1 = np.array([0, 1, 1, 5], dtype=np.int32)
+        splits_val2 = np.array([0, 3, 3, 5, 9, 10], dtype=np.int32)
+        dense_vals_val = np.array([10, 11, 12, 13, 14, 15, 16, 17, 18, 19], dtype=np.float32)
+        def func(splits1, splits2, rt_dense_values):
+            x = tf.RaggedTensor.from_nested_row_splits(rt_dense_values, [splits1, splits2], validate=True)
+            y = x.to_tensor(default_value=7)
+            return tf.identity(y, name=_TFOUTPUT)
+        self._run_test_case(func, [_OUTPUT], {_INPUT: splits_val1, _INPUT1: splits_val2, _INPUT2: dense_vals_val})
+
+    @check_tf_min_version("2.2", "ragged to_tensor with constrained shape")
+    @check_opset_min_version(11, "CumSum")
+    def test_ragged_tensor_to_tensor_constrain_shape(self):
+        splits_val1 = np.array([0, 1, 1, 5], dtype=np.int32)
+        splits_val2 = np.array([0, 3, 3, 5, 9, 10], dtype=np.int32)
+        dense_vals_val = np.array([10, 11, 12, 13, 14, 15, 16, 17, 18, 19], dtype=np.float32)
+        def func(splits1, splits2, rt_dense_values):
+            x = tf.RaggedTensor.from_nested_row_splits(rt_dense_values, [splits1, splits2], validate=True)
+            y = x.to_tensor(default_value=7, shape=[20, None, 2])
+            return tf.identity(y, name=_TFOUTPUT)
+        self._run_test_case(func, [_OUTPUT], {_INPUT: splits_val1, _INPUT1: splits_val2, _INPUT2: dense_vals_val})
+
     @check_tf_min_version("1.14", "ragged needs tf 1.14")
     @check_opset_min_version(11, "Range")
     def test_ragged_range_float(self):

tf2onnx/onnx_opset/tensor.py

@@ -423,7 +423,10 @@ class GatherV2:
     @classmethod
     def version_1(cls, ctx, node, **kwargs):
         # for GatherV2 axis come as input
+        err_msg = "Opset 12 required for batch_dims attribute of GatherV2"
+        utils.make_sure(node.get_attr_value("batch_dims", 0) == 0, err_msg)
         node.type = "Gather"
+        utils.make_sure(node.inputs[2].is_const(), "Axis of GatherV2 node must be constant")
         axis = node.inputs[2].get_tensor_value()
         ctx.remove_input(node, node.input[2], 2)
         node.set_attr("axis", axis)
@@ -433,6 +436,42 @@ def version_11(cls, ctx, node, **kwargs):
         # no change
         cls.version_1(ctx, node, **kwargs)
 
+    @classmethod
+    def version_12(cls, ctx, node, **kwargs):
+        batch_dims = node.get_attr_value("batch_dims", 0)
+        if batch_dims == 0:
+            cls.version_1(ctx, node, **kwargs)
+            return
+        # If batch_dims is not zero, use GatherND to simulate Gather with batch dims.
+        data_inp, indices_inp, axis_inp = node.input
+        utils.make_sure(node.inputs[2].is_const(), "Axis of GatherV2 node must be constant")
+        axis = node.inputs[2].get_tensor_value()
+        ctx.remove_input(node, axis_inp, 2)
+        if ctx.get_dtype(indices_inp) != TensorProto.INT64:
+            indices_inp = ctx.make_node("Cast", [indices_inp], attr={'to': TensorProto.INT64}).output[0]
+        unperm = None
+        # GatherND doesn't take an axis so we have to transpose stuff around
+        if axis != batch_dims:
+            data_rank = ctx.get_rank(data_inp)
+            indices_rank = ctx.get_rank(indices_inp)
+            result_rank = data_rank + indices_rank - 1 - batch_dims
+            shift_amt = axis - batch_dims
+            err_msg = "Cannot convert GatherV2 with batch dims since inputs have unknown ranks."
+            utils.make_sure(data_rank is not None and indices_rank is not None, err_msg)
+            perm = list(range(data_rank))
+            perm = perm[:batch_dims] + perm[axis:axis+1] + perm[batch_dims:axis] + perm[axis+1:]
+            data_inp = ctx.make_node("Transpose", [data_inp], attr={'perm': perm}).output[0]
+            ctx.replace_input(node, node.input[0], data_inp, 0)
+            unperm = list(range(result_rank))
+            j = indices_rank+shift_amt
+            unperm = unperm[:batch_dims] + unperm[indices_rank:j] + unperm[batch_dims:indices_rank] + unperm[j:]
+        node.type = "GatherND"
+        unsqueeze_node = GraphBuilder(ctx).make_unsqueeze({'data': indices_inp, 'axes': [-1]})
+        ctx.replace_input(node, node.input[1], unsqueeze_node, 1)
+        if unperm is not None:
+            ctx.update_node_shape_dtype(node, override=True)
+            ctx.insert_new_node_on_output("Transpose", node.output[0], perm=unperm)
+
 
 def _make_gathernd_inner_loop(ctx, params, index, dtype):
     """create the inner loop for GatherNd."""
@@ -2077,43 +2116,77 @@ def ragged_lengths_to_sparse_indices(ctx, ragged_lens):
     return num_rows, num_cols, row_indices, col_indices
 
 
+def ragged_nested_splits_to_sparse_indices(ctx, nested_splits, op_name_scope):
+    sparse_indices = None
+    dense_shape_dims = []
+    for split in nested_splits:
+        if ctx.get_dtype(split) != TensorProto.INT64:
+            split = ctx.make_node("Cast", [split], attr={'to': TensorProto.INT64}).output[0]
+        max_int64 = int(utils.get_max_value(np.int64))
+        slice1 = GraphBuilder(ctx).make_slice(
+            {"data": split, "ends": [max_int64], "starts": [1], "axes": [0]})
+        slice2 = GraphBuilder(ctx).make_slice(
+            {"data": split, "ends": [-1], "starts": [0], "axes": [0]})
+        ragged_lens = ctx.make_node("Sub", [slice1, slice2]).output[0]
+        num_rows, num_cols, row_indices, col_indices = ragged_lengths_to_sparse_indices(ctx, ragged_lens)
+        if not dense_shape_dims:
+            dense_shape_dims.append(num_rows)
+        dense_shape_dims.append(num_cols)
+        if sparse_indices is None:
+            row_indices = GraphBuilder(ctx).make_unsqueeze({"data": row_indices, "axes": [1]})
+        else:
+            row_indices = ctx.make_node("Gather", [sparse_indices, row_indices]).output[0]
+        col_indices = GraphBuilder(ctx).make_unsqueeze({"data": col_indices, "axes": [1]})
+        sparse_indices = ctx.make_node("Concat", [row_indices, col_indices], attr={'axis': 1},
+                                       op_name_scope=op_name_scope).output[0]
+    dense_shape = ctx.make_node("Concat", dense_shape_dims, attr={'axis': 0}, op_name_scope=op_name_scope).output[0]
+    return sparse_indices, dense_shape
+
+
 @tf_op("RaggedTensorToSparse")
 class RaggedTensorToSparse:
     @classmethod
     def version_11(cls, ctx, node, **kwargs):
         # https://www.tensorflow.org/guide/ragged_tensor#multiple_ragged_dimensions
         dense_values = node.input[-1]
         nested_splits = node.input[:-1]
-        sparse_indices = None
-        dense_shape_dims = []
-        for split in nested_splits:
-            if ctx.get_dtype(split) != TensorProto.INT64:
-                split = ctx.make_node("Cast", [split], attr={'to': TensorProto.INT64}).output[0]
-            max_int64 = int(utils.get_max_value(np.int64))
-            slice1 = GraphBuilder(ctx).make_slice(
-                {"data": split, "ends": [max_int64], "starts": [1], "axes": [0]})
-            slice2 = GraphBuilder(ctx).make_slice(
-                {"data": split, "ends": [-1], "starts": [0], "axes": [0]})
-            ragged_lens = ctx.make_node("Sub", [slice1, slice2]).output[0]
-            num_rows, num_cols, row_indices, col_indices = ragged_lengths_to_sparse_indices(ctx, ragged_lens)
-            if not dense_shape_dims:
-                dense_shape_dims.append(num_rows)
-            dense_shape_dims.append(num_cols)
-            if sparse_indices is None:
-                row_indices = GraphBuilder(ctx).make_unsqueeze({"data": row_indices, "axes": [1]})
-            else:
-                row_indices = ctx.make_node("Gather", [sparse_indices, row_indices]).output[0]
-            col_indices = GraphBuilder(ctx).make_unsqueeze({"data": col_indices, "axes": [1]})
-            sparse_indices = ctx.make_node("Concat", [row_indices, col_indices], attr={'axis': 1},
-                                           op_name_scope=node.name).output[0]
-        dense_shape = ctx.make_node("Concat", dense_shape_dims, attr={'axis': 0}, op_name_scope=node.name).output[0]
-
+        sparse_indices, dense_shape = ragged_nested_splits_to_sparse_indices(ctx, nested_splits, node.name)
         ctx.replace_all_inputs(node.output[0], sparse_indices)
         ctx.replace_all_inputs(node.output[1], dense_values)
         ctx.replace_all_inputs(node.output[2], dense_shape)
         ctx.remove_node(node.name)
 
 
+@tf_op("RaggedTensorToTensor")
+class RaggedTensorToTensor:
+    @classmethod
+    def version_11(cls, ctx, node, **kwargs):
+        shape, values, default_value, *row_partition_tensors = node.input
+        partition_types = node.get_attr_value("row_partition_types")
+        error_msg = "Only ROW_SPLITS partition type is supported for RaggedTensorToTensor. types: %r"
+        utils.make_sure(all(t == b'ROW_SPLITS' for t in partition_types), error_msg, partition_types)
+        nested_splits = row_partition_tensors
+        sparse_indices, dense_shape = ragged_nested_splits_to_sparse_indices(ctx, nested_splits, node.name)
+        # A shape of rank 0 means the natural shape should be used.
+        if ctx.get_rank(shape) != 0:
+            if ctx.get_dtype(shape) != TensorProto.INT64:
+                shape = ctx.make_node("Cast", [shape], attr={'to': TensorProto.INT64}).output[0]
+            const_zero_int64 = ctx.make_const(utils.make_name("const_zero"), np.array(0, dtype=np.int64)).output[0]
+            unspec_dims = ctx.make_node("Less", [shape, const_zero_int64]).output[0]
+            out_shape = ctx.make_node("Where", [unspec_dims, dense_shape, shape]).output[0]
+            out_shape_unsq = GraphBuilder(ctx).make_unsqueeze({'data': out_shape, 'axes': [0]})
+            amt_idx_in_bounds = ctx.make_node("Sub", [out_shape_unsq, sparse_indices]).output[0]
+            amt_in_bounds_flat = ctx.make_node("ReduceMin", [amt_idx_in_bounds], attr={'axes': [1], 'keepdims': False})
+            idx_in_bounds = ctx.make_node("Greater", [amt_in_bounds_flat.output[0], const_zero_int64]).output[0]
+            sparse_indices = ctx.make_node("Compress", [sparse_indices, idx_in_bounds], attr={'axis': 0}).output[0]
+            values = ctx.make_node("Compress", [values, idx_in_bounds], attr={'axis': 0}).output[0]
+        else:
+            out_shape = dense_shape
+        expand_node = ctx.make_node("Expand", [default_value, out_shape])
+        node.type = "ScatterND"
+        ctx.replace_inputs(node, [expand_node.output[0], sparse_indices, values])
+
+
 @tf_op("RaggedRange")
 class RaggedRange:
     @classmethod