Implement conversion of RaggedToVariant and RaggedFromVariant in loops (#1503)

Signed-off-by: Tom Wildenhain <tomwi@microsoft.com>

Author: TomWildenhain-Microsoft

tests/backend_test_base.py

@@ -272,6 +272,16 @@ def get_shape(info):
             if not info.type.tensor_type.HasField("shape"):
                 return None
             return [d.dim_value if d.HasField('dim_value') else -1 for d in info.type.tensor_type.shape.dim]
+        def get_dtype(info):
+            tensor_type = info.type.tensor_type
+            is_seq = False
+            result = None
+            if info.type.HasField("sequence_type"):
+                tensor_type = info.type.sequence_type.elem_type.tensor_type
+                is_seq = True
+            if tensor_type.HasField("elem_type"):
+                result = tensor_type.elem_type
+            return utils.SeqType(result) if is_seq else result
         for info in model_shapes.graph.value_info:
             if info.name == "":
                 continue
@@ -289,7 +299,7 @@ def get_shape(info):
                         self.assertEqual(d1, d2)
             else:
                 self.assertEqual(onnx_shape, tf2onnx_shape)
-            self.assertEqual(info.type.tensor_type.elem_type, graph.get_dtype(info.name))
+            self.assertEqual(get_dtype(info), graph.get_dtype(info.name))
 
     def run_test_case(self, func, feed_dict, input_names_with_port, output_names_with_port,
                       rtol=1e-07, atol=1e-5, mtol=None, convert_var_to_const=True, constant_fold=True,

tests/test_backend.py

@@ -4581,6 +4581,42 @@ def func(starts, limits, deltas):
         self._run_test_case(func, [_OUTPUT, _OUTPUT1], {_INPUT: starts_val, _INPUT1: limits_val,
                                                         _INPUT2: deltas_val})
 
+    @check_tf_min_version("2.0", "ragged variant needs tf 2.0")
+    @check_opset_min_version(13, "Loop over tensor sequences")
+    def test_ragged_to_variant(self):
+        splits_val = np.array([0, 3, 3, 5, 9, 10], dtype=np.int32)
+        dense_vals_val = np.arange(10 * 3 * 2, dtype=np.float32).reshape([10, 3, 2])
+
+        def fn(elem):
+            res = elem + elem * elem
+            return res
+
+        def func(splits, rt_dense_values):
+            x = tf.RaggedTensor.from_nested_row_splits(rt_dense_values, [splits], validate=True)
+            y = tf.map_fn(fn, x)
+            return tf.identity(y.row_splits, name=_TFOUTPUT), tf.identity(y.flat_values, name=_TFOUTPUT1)
+        self._run_test_case(func, [_OUTPUT, _OUTPUT1], {_INPUT: splits_val, _INPUT1: dense_vals_val})
+
+    @check_tf_min_version("2.0", "ragged variant needs tf 2.0")
+    @check_opset_min_version(13, "Loop over tensor sequences")
+    def test_ragged_to_variant_unknown_shape(self):
+        splits_val = np.array([0, 3, 3, 5, 9, 10], dtype=np.int64)
+        dense_vals_shape = np.array([10, 3, 2], dtype=np.int32)
+        splits_pads_val = np.array([[0, 0]], dtype=np.int32)
+
+        def fn(elem):
+            res = elem + elem * elem
+            return res
+
+        def func(splits, rt_dense_values_shape, splits_pads):
+            rt_dense_values = tf.ones(rt_dense_values_shape, dtype=tf.int32)
+            splits = tf.pad(splits, splits_pads)
+            x = tf.RaggedTensor.from_nested_row_splits(rt_dense_values, [splits], validate=True)
+            y = tf.map_fn(fn, x)
+            return tf.identity(y.row_splits, name=_TFOUTPUT), tf.identity(y.flat_values, name=_TFOUTPUT1)
+        self._run_test_case(func, [_OUTPUT, _OUTPUT1],
+                            {_INPUT: splits_val, _INPUT1: dense_vals_shape, _INPUT2: splits_pads_val})
+
     @check_opset_min_version(9, "Compress")
     def test_dynamic_partition_both_vector(self):
         data_val = np.array([1, 2, 3, 4, 5, 6, 7, 8], dtype=np.float32)

tf2onnx/graph.py

@@ -469,6 +469,8 @@ def __init__(self, nodes, output_shapes=None, dtypes=None, target=None, opset=No
         # Used by the tflite while loop handler
         self.scan_outputs = []
         self.func_inputs = []
+        self.ragged_variant_list_reads = []
+        self.ragged_variant_list_writes = []
 
         self._target = set(target)
         self._dtypes = dtypes

tf2onnx/onnx_opset/controlflow.py

@@ -19,6 +19,7 @@
 from tf2onnx import utils
 from tf2onnx.handler import tf_op
 from tf2onnx.tf_loader import find_function
+from tf2onnx.graph_builder import GraphBuilder
 
 
 logger = logging.getLogger(__name__)
@@ -401,6 +402,7 @@ def version_7(cls, ctx, node, **kwargs):
         cond_input_to_state_var = {}
         scan_outputs = []
         input_idx_to_remove = []
+        idx_to_ragged_writes = dict(body.ragged_variant_list_writes)
         # remove TensorListReserve
         for idx, name in enumerate(tf_while_inputs):
             if idx == 1:
@@ -416,9 +418,15 @@ def version_7(cls, ctx, node, **kwargs):
                 # there is no equivalent step in onnx and we should remove it.
                 output_shape = None
                 output_dtype = n.get_attr_value("element_dtype")
+                is_ragged = False
                 if n.type == "TensorListReserve" and n.inputs[0].is_const() and not n.inputs[0].is_scalar():
                     output_shape = [-1] + n.inputs[0].get_tensor_value(as_list=True)
-                scan_outputs.append((idx, n, output_shape, output_dtype))
+                if idx in idx_to_ragged_writes:
+                    output_shape = None
+                    output_dtype = body.get_dtype(idx_to_ragged_writes[idx].input[0])
+                    is_ragged = True
+                    loop_vars.append(name)
+                scan_outputs.append((idx, n, output_shape, output_dtype, is_ragged))
                 continue
 
             # tensor arrays we read from can't be loop_vars and we fetch them from the outer context instead
@@ -437,8 +445,29 @@ def version_7(cls, ctx, node, **kwargs):
             del body.outputs[idx]
 
         scan_output_names = []
-        # remove tensor array that are passed in to the loop
-        for idx, n, output_shape, output_dtype in reversed(scan_outputs):
+        ragged_scan_output_names = []
+        ragged_scan_output_to_len = {}
+
+        # remove tensor arrays that are passed in to the loop
+        for idx, n, output_shape, output_dtype, is_ragged in reversed(scan_outputs):
+            if is_ragged:
+                out = n.output[0]
+                ctx.remove_node(n.name)
+                seq_empty = ctx.make_node("SequenceEmpty", [], attr={'dtype': output_dtype}, name=n.name,
+                                          outputs=[out], shapes=[None], dtypes=[utils.SeqType(output_dtype)])
+                ctx.replace_all_inputs(n.output[0], seq_empty.output[0])
+                # Ragged tensors also must track the length of each row
+                output_shapes.append([-1])
+                output_dtypes.append(TensorProto.INT64)
+                output_shapes[idx] = None
+                output_dtypes[idx] = utils.SeqType(output_dtype)
+                body_ragged_name = utils.make_name("ragged_scan_output")
+                external_ragged_name = utils.make_name("ragged_output")
+                scan_output_names.append(body_ragged_name)
+                output_names.append(external_ragged_name)
+                ragged_scan_output_names.append(body_ragged_name)
+                ragged_scan_output_to_len[output_names[idx]] = external_ragged_name
+                continue
             ctx.remove_node(n.name)
             # make the node output bad
             ctx.replace_all_inputs(n.output[0], "@@ALLOC")  # ops=ctx.get_nodes()
@@ -475,11 +504,16 @@ def version_7(cls, ctx, node, **kwargs):
 
         # shift output consumers
         for k, v in output_map.items():
-            ctx.replace_all_inputs(k, v)  # ops=ctx.get_nodes()
+            if k not in ragged_scan_output_to_len.values():
+                ctx.replace_all_inputs(k, v)  # ops=ctx.get_nodes()
+
+        ragged_scan_output_to_len = {output_map[k]: output_map[v] for k, v in ragged_scan_output_to_len.items()}
 
         wire_while_body(ctx, body, loop_node, body_input_to_state_var, cond_input_to_state_var, output_shapes,
-                        output_dtypes, body_name, node.name, cond_graph, tf_while_inputs, scan_output_names)
+                        output_dtypes, body_name, node.name, cond_graph, tf_while_inputs, scan_output_names,
+                        ragged_scan_output_names)
 
+        loop_node.ragged_scan_output_to_len = ragged_scan_output_to_len
         # if there was a tensorflow variant type, bind in a real type here
         # FIXME: I don't think this is needed anymore
         for i, n in enumerate(body.inputs):
@@ -488,7 +522,8 @@ def version_7(cls, ctx, node, **kwargs):
 
 
 def wire_while_body(parent_g, g, loop_node, body_input_to_state_var, cond_input_to_state_var, output_shapes,
-                    output_dtypes, scope, parent, cond_graph, tf_while_inputs, scan_output_names):
+                    output_dtypes, scope, parent, cond_graph, tf_while_inputs, scan_output_names,
+                    ragged_scan_output_names):
     """Wire subgraph graph into main."""
     remove_parents = []
     to_remove = []
@@ -519,8 +554,25 @@ def wire_while_body(parent_g, g, loop_node, body_input_to_state_var, cond_input_
 
     # this is a tensor array write - make it an identity
     scan_outputs = []
+    ragged_scan_outputs_cnt = 0
+    names_to_scan_outputs = {}
+
     for node in g.get_nodes():
         if node.type == "TensorListSetItem":
+            if node.inputs[2].type == "RaggedTensorToVariant":
+                node.type = "SequenceInsert"
+                row_content = node.inputs[2].input[0]
+                g.replace_inputs(node, [node.input[0], row_content])
+                g.set_shape(node.output[0], g.get_shape(node.input[1]))
+                g.set_dtype(node.output[0], utils.SeqType(g.get_dtype(node.input[1])))
+                dense_shape = g.make_node("Shape", [row_content]).output[0]
+                zero_const = g.make_const(utils.make_name("zero_const"), np.array(0, np.int64)).output[0]
+                row_length = g.make_node("Gather", [dense_shape, zero_const]).output[0]
+                row_length_id = g.make_node("Identity", [row_length])
+                scan_outputs.append(row_length_id.output[0])
+                names_to_scan_outputs[ragged_scan_output_names[ragged_scan_outputs_cnt]] = row_length_id.output[0]
+                ragged_scan_outputs_cnt += 1
+                continue
             remove_parents.append(node.input[0])
             node.type = "Identity"
             g.set_shape(node.output[0], g.get_shape(node.input[2]))
@@ -531,8 +583,9 @@ def wire_while_body(parent_g, g, loop_node, body_input_to_state_var, cond_input_
     if len(scan_outputs) != len(scan_output_names):
         raise ValueError("While loop couldn't find scan output index for nodes")
 
-    names_to_scan_outputs = {}
     for output in scan_outputs:
+        if output in names_to_scan_outputs.values():
+            continue
         last_output = output
         consumers = g.find_output_consumers(last_output)
         while consumers:
@@ -547,8 +600,9 @@ def wire_while_body(parent_g, g, loop_node, body_input_to_state_var, cond_input_
 
     # Reorder scan outputs
     scan_outputs = [names_to_scan_outputs[name] for name in scan_output_names]
+
+    # Use shapes from subgraph if loop node shapes for scan outputs are missing
     for i in range(-len(scan_output_names), 0):
-        # Use shapes from subgraph if loop node shapes for scan outputs are missing
         if loop_node.output_shapes[i] is None:
             shape = g.get_shape(scan_outputs[i])
             if shape is not None:
@@ -580,6 +634,31 @@ def wire_while_body(parent_g, g, loop_node, body_input_to_state_var, cond_input_
             if node.type in ["Identity"]:
                 g.set_dtype(o, node.inputs[0].output_dtypes[0])
 
+    for node in g.ragged_variant_list_reads:
+        # Requires opset 11
+        gather = node.inputs[0]
+        inp = gather.inputs[0]
+        while inp.type == "Identity":
+            inp = inp.inputs[0]
+        err_msg1 = "Could not find corresponding RaggedTensorToVariant for node %s" % node.name
+        err_msg2 = "Input to RaggedTensorToVariant for loop has batched_input=False for node %s" % inp.name
+        err_msg3 = "RAGGED_RANK != 1 for RaggedTensorToVariant node %s" % node.name
+        utils.make_sure(inp.type == "RaggedTensorToVariant", err_msg1)
+        utils.make_sure(inp.get_attr_value("batched_input"), err_msg2)
+        utils.make_sure(inp.get_attr_value("RAGGED_RANK") == 1, err_msg3)
+        idx = gather.input[1]
+        idx_unsq = GraphBuilder(g).make_unsqueeze({'data': idx, 'axes': [0]})
+        np_dtype = utils.map_onnx_to_numpy_type(g.get_dtype(idx_unsq))
+        const_one = g.make_const(utils.make_name("const_1"), np.array(1, np_dtype)).output[0]
+        idx_plus_1 = g.make_node("Add", [idx_unsq, const_one]).output[0]
+        splits, values = inp.input
+        start = g.make_node("Gather", [splits, idx_unsq]).output[0]
+        end = g.make_node("Gather", [splits, idx_plus_1]).output[0]
+        np_dtype2 = utils.map_onnx_to_numpy_type(g.get_dtype(splits))
+        axes = g.make_const(utils.make_name("const_zero"), np.array([0], np_dtype2)).output[0]
+        sliced_vals = g.make_node("Slice", [values, start, end, axes]).output[0]
+        g.replace_all_inputs(node.output[0], sliced_vals)
+
     return g
 
 

tf2onnx/onnx_opset/tensor.py

@@ -2584,6 +2584,63 @@ def version_11(cls, ctx, node, **kwargs):
         ctx.remove_node(node.name)
 
 
+@tf_op("RaggedTensorFromVariant")
+class RaggedTensorFromVariant:
+    @classmethod
+    def version_13(cls, ctx, node, **kwargs):
+        inp = node.inputs[0]
+        if inp.is_while():
+            row_lengths = inp.ragged_scan_output_to_len.get(node.input[0])
+            utils.make_sure(row_lengths is not None, "Couldn't find lengths for %s node %s" % (node.type, node.name))
+            dense_values = ctx.make_node("ConcatFromSequence", [node.input[0]], attr={'axis': 0}).output[0]
+            const_zero = ctx.make_const(utils.make_name("const_zero"), np.array(0, np.int64)).output[0]
+            const_zero_unsq = ctx.make_const(utils.make_name("const_zero"), np.array([0], np.int64)).output[0]
+            row_splits = ctx.make_node("CumSum", [row_lengths, const_zero]).output[0]
+            row_splits_w_zero = ctx.make_node("Concat", [const_zero_unsq, row_splits], attr={'axis': 0}).output[0]
+            idx_dtype = ctx.get_dtype(node.output[0])
+            if idx_dtype != TensorProto.INT64:
+                row_splits_w_zero = ctx.make_node("Cast", [row_splits_w_zero], attr={'to': idx_dtype}).output[0]
+            ctx.replace_all_inputs(node.output[0], row_splits_w_zero)
+            ctx.replace_all_inputs(node.output[1], dense_values)
+            ctx.remove_node(node.name)
+            return
+
+        utils.make_sure(inp.type == "Gather", "RaggedTensorFromVariant only supported after TensorListGetItem")
+        variant = inp.inputs[0]
+        err_msg = "RaggedTensorFromVariant only supported if variant is a graph input"
+        # Variant input will be found during loop conversion
+        utils.make_sure(variant.type == "Placeholder", err_msg)
+        ctx.ragged_variant_list_reads.append(node)
+
+
+@tf_op("RaggedTensorToVariant")
+class RaggedTensorToVariant:
+    @classmethod
+    def version_13(cls, ctx, node, **kwargs):
+        cons = ctx.find_output_consumers(node.output[0])
+        err_msg = "RaggedTensorToVariant only supported as input/output to loops"
+        utils.make_sure(len(cons) == 1, err_msg)
+        if cons[0].type == "TensorListFromTensor":
+            # Will be delt with in loop
+            cons = ctx.find_output_consumers(cons[0].output[0])
+            utils.make_sure(all(n.is_while() for n in cons), err_msg)
+            return
+        utils.make_sure(cons[0].type == "TensorListSetItem", err_msg)
+        tensor_set_item = cons[0]
+        list_output = tensor_set_item.output[0]
+        cons = ctx.find_output_consumers(list_output)
+        while len(cons) == 1 and cons[0].type == "Identity":
+            list_output = cons[0].output[0]
+            cons = ctx.find_output_consumers(list_output)
+        utils.make_sure(not cons, err_msg)
+        utils.make_sure(list_output in ctx.outputs, err_msg)
+        err_msg2 = "RaggedTensorToVariant within loop requires RAGGED_RANK=0"
+        err_msg3 = "RaggedTensorToVariant within loop requires batched_input=False"
+        utils.make_sure(node.get_attr_value("RAGGED_RANK") == 0, err_msg2)
+        utils.make_sure(not node.get_attr_value("batched_input"), err_msg3)
+        ctx.ragged_variant_list_writes.append((ctx.outputs.index(list_output), node))
+
+
 @tf_op("SparseReshape")
 class SparseReshape:
     @classmethod

tf2onnx/rewriter/__init__.py

@@ -24,6 +24,7 @@
 from tf2onnx.rewriter.conv2d_with_add_rewriter import rewrite_biasadd_with_conv2d
 from tf2onnx.rewriter.quantization_ops_rewriter import rewrite_quantize_and_dequantize
 from tf2onnx.rewriter.layer_normalization_rewriter import rewrite_layer_normalization
+from tf2onnx.rewriter.ragged_variant_shape_rewriter import rewrite_ragged_variant_shape
 
 
 __all__ = [
@@ -48,5 +49,6 @@
     "rewrite_biasadd_with_conv2d",
     "rewrite_quantize_and_dequantize",
     "rewrite_layer_normalization",
-    "rewrite_conv_dilations"
+    "rewrite_conv_dilations",
+    "rewrite_ragged_variant_shape"
 ]

tf2onnx/rewriter/ragged_variant_shape_rewriter.py

@@ -0,0 +1,39 @@
+# SPDX-License-Identifier: Apache-2.0
+
+
+"""
+tf2onnx.rewriter - RaggedTensorToVariant -> Shape pattern
+"""
+
+import numpy as np
+from tf2onnx import utils
+from tf2onnx.graph_matcher import OpTypePattern, GraphMatcher
+
+
+# pylint: disable=missing-docstring
+
+
+def rewrite_ragged_variant_shape(g, ops):
+    pattern1 = \
+        OpTypePattern('Shape', name='shape', inputs=[
+            OpTypePattern('RaggedTensorToVariant', name='raggedtovariant')
+        ])
+
+    pattern_list = [pattern1]
+    for pattern in pattern_list:
+        matcher = GraphMatcher(pattern)
+        match_results = list(matcher.match_ops(ops))
+        for match in match_results:
+            shape = match.get_op('shape')
+            raggedtovariant = match.get_op('raggedtovariant')
+            if raggedtovariant.get_attr_value("batched_input") != 1:
+                continue
+            if raggedtovariant.get_attr_value("RAGGED_RANK") != 1:
+                continue
+            # Shape of batched variant from ragged is same as number of splits minus 1
+            g.replace_inputs(shape, [raggedtovariant.input[0]])
+            np_dtype = utils.map_onnx_to_numpy_type(g.get_dtype(shape.output[0]))
+            const_one = g.make_const(utils.make_name("const_one"), np.array(1, np_dtype)).output[0]
+            g.insert_new_node_on_output("Sub", shape.output[0], inputs=[shape.output[0], const_one])
+
+    return ops

tf2onnx/tfonnx.py

@@ -619,6 +619,7 @@ def compat_handler(ctx, node, **kwargs):
         rewrite_biasadd_with_conv2d,
         rewrite_layer_normalization,
         rewrite_gemm,
+        rewrite_ragged_variant_shape,
     ]
 
     if custom_rewriter is not None: