Handles sequences

xadupre · xadupre · commit db3b67c6c33d · 2025-06-03T15:59:14.000+02:00
diff --git a/_doc/api/reference/torch_ops/control_flow_ops.rst b/_doc/api/reference/torch_ops/control_flow_ops.rst
@@ -0,0 +1,6 @@
+
+onnx_diagnostic.reference.torch_ops.control_flow_ops
+====================================================
+
+.. automodule:: onnx_diagnostic.reference.torch_ops.control_flow_ops
+    :members:
diff --git a/_doc/api/reference/torch_ops/generator_ops.rst b/_doc/api/reference/torch_ops/generator_ops.rst
@@ -0,0 +1,6 @@
+
+onnx_diagnostic.reference.torch_ops.generator_ops
+=================================================
+
+.. automodule:: onnx_diagnostic.reference.torch_ops.generator_ops
+    :members:
diff --git a/_doc/api/reference/torch_ops/index.rst b/_doc/api/reference/torch_ops/index.rst
@@ -9,6 +9,14 @@ onnx_diagnostic.reference.torch_ops
 
     access_ops
     binary_ops
+    control_flow_ops
+    generator_ops
+    nn_ops
+    other_ops
+    reduce_ops
+    sequence_ops
+    shape_ops
+    unary_ops
 
 OpRun
 +++++
@@ -22,6 +30,12 @@ OpRunValue
 .. autoclass:: onnx_diagnostic.reference.torch_ops.OpRunValue
     :members:
 
+OpRunValueSequence
+++++++++++++++++++
+
+.. autoclass:: onnx_diagnostic.reference.torch_ops.OpRunValueSequence
+    :members:
+
 OpRunFunction
 +++++++++++++
 
diff --git a/_doc/api/reference/torch_ops/nn_ops.rst b/_doc/api/reference/torch_ops/nn_ops.rst
@@ -0,0 +1,6 @@
+
+onnx_diagnostic.reference.torch_ops.nn_ops
+==========================================
+
+.. automodule:: onnx_diagnostic.reference.torch_ops.nn_ops
+    :members:
diff --git a/_doc/api/reference/torch_ops/other_ops.rst b/_doc/api/reference/torch_ops/other_ops.rst
@@ -0,0 +1,6 @@
+
+onnx_diagnostic.reference.torch_ops.other_ops
+=============================================
+
+.. automodule:: onnx_diagnostic.reference.torch_ops.other_ops
+    :members:
diff --git a/_doc/api/reference/torch_ops/reduce_ops.rst b/_doc/api/reference/torch_ops/reduce_ops.rst
@@ -0,0 +1,6 @@
+
+onnx_diagnostic.reference.torch_ops.reduce_ops
+==============================================
+
+.. automodule:: onnx_diagnostic.reference.torch_ops.reduce_ops
+    :members:
diff --git a/_doc/api/reference/torch_ops/sequence_ops.rst b/_doc/api/reference/torch_ops/sequence_ops.rst
@@ -0,0 +1,6 @@
+
+onnx_diagnostic.reference.torch_ops.sequence_ops
+================================================
+
+.. automodule:: onnx_diagnostic.reference.torch_ops.sequence_ops
+    :members:
diff --git a/_doc/api/reference/torch_ops/shape_ops.rst b/_doc/api/reference/torch_ops/shape_ops.rst
@@ -0,0 +1,6 @@
+
+onnx_diagnostic.reference.torch_ops.shape_ops
+=============================================
+
+.. automodule:: onnx_diagnostic.reference.torch_ops.shape_ops
+    :members:
diff --git a/_doc/api/reference/torch_ops/unary_ops.rst b/_doc/api/reference/torch_ops/unary_ops.rst
@@ -0,0 +1,6 @@
+
+onnx_diagnostic.reference.torch_ops.unary_ops
+=============================================
+
+.. automodule:: onnx_diagnostic.reference.torch_ops.unary_ops
+    :members:
diff --git a/_unittests/ut_reference/test_torch_onnx_evaluator.py b/_unittests/ut_reference/test_torch_onnx_evaluator.py
@@ -956,6 +956,119 @@ def _mkv_(name):
             torch.tensor([1], dtype=torch.float32),
         )
 
+    def test_loop(self):
+        x = np.array([1, 2, 3, 4, 5]).astype(np.float32)
+
+        model = oh.make_model(
+            graph=oh.make_graph(
+                name="loop_test",
+                inputs=[
+                    oh.make_tensor_value_info("trip_count", TINT64, ["a"]),
+                    oh.make_tensor_value_info("cond", onnx.TensorProto.BOOL, [1]),
+                ],
+                outputs=[oh.make_tensor_value_info("res", TFLOAT, [])],
+                nodes=[
+                    oh.make_node("SequenceEmpty", [], ["seq_empty"], dtype=TFLOAT),
+                    oh.make_node(
+                        "Loop",
+                        inputs=["trip_count", "cond", "seq_empty"],
+                        outputs=["seq_res"],
+                        body=oh.make_graph(
+                            [
+                                oh.make_node(
+                                    "Identity", inputs=["cond_in"], outputs=["cond_out"]
+                                ),
+                                oh.make_node(
+                                    "Constant",
+                                    inputs=[],
+                                    outputs=["x"],
+                                    value=oh.make_tensor(
+                                        name="const_tensor_x",
+                                        data_type=TFLOAT,
+                                        dims=x.shape,
+                                        vals=x.flatten().astype(float),
+                                    ),
+                                ),
+                                oh.make_node(
+                                    "Constant",
+                                    inputs=[],
+                                    outputs=["one"],
+                                    value=oh.make_tensor(
+                                        name="const_tensor_one",
+                                        data_type=TINT64,
+                                        dims=(),
+                                        vals=[1],
+                                    ),
+                                ),
+                                oh.make_node(
+                                    "Constant",
+                                    inputs=[],
+                                    outputs=["slice_start"],
+                                    value=oh.make_tensor(
+                                        name="const_tensor_zero",
+                                        data_type=TINT64,
+                                        dims=(1,),
+                                        vals=[0],
+                                    ),
+                                ),
+                                oh.make_node(
+                                    "Add", inputs=["iter_count", "one"], outputs=["end"]
+                                ),
+                                oh.make_node(
+                                    "Constant",
+                                    inputs=[],
+                                    outputs=["axes"],
+                                    value=oh.make_tensor(
+                                        name="const_tensor_axes",
+                                        data_type=TINT64,
+                                        dims=(1,),
+                                        vals=[0],
+                                    ),
+                                ),
+                                oh.make_node(
+                                    "Unsqueeze", inputs=["end", "axes"], outputs=["slice_end"]
+                                ),
+                                oh.make_node(
+                                    "Slice",
+                                    inputs=["x", "slice_start", "slice_end"],
+                                    outputs=["slice_out"],
+                                ),
+                                oh.make_node(
+                                    "SequenceInsert",
+                                    inputs=["seq_in", "slice_out"],
+                                    outputs=["seq_out"],
+                                ),
+                            ],
+                            "loop_body",
+                            [
+                                oh.make_tensor_value_info("iter_count", TINT64, []),
+                                oh.make_tensor_value_info(
+                                    "cond_in", onnx.TensorProto.BOOL, []
+                                ),
+                                oh.make_tensor_sequence_value_info("seq_in", TFLOAT, None),
+                            ],
+                            [
+                                oh.make_tensor_value_info(
+                                    "cond_out", onnx.TensorProto.BOOL, []
+                                ),
+                                oh.make_tensor_sequence_value_info("seq_out", TFLOAT, None),
+                            ],
+                        ),
+                    ),
+                    oh.make_node(
+                        "ConcatFromSequence",
+                        inputs=["seq_res"],
+                        outputs=["res"],
+                        axis=0,
+                        new_axis=0,
+                    ),
+                ],
+            )
+        )
+        self._finalize_test(
+            model, torch.tensor(5, dtype=torch.int64), torch.tensor(1, dtype=torch.bool)
+        )
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)
diff --git a/onnx_diagnostic/reference/torch_evaluator.py b/onnx_diagnostic/reference/torch_evaluator.py
@@ -383,7 +383,9 @@ def run_with_values(
         # inputs
         for k, v in zip(self.input_names, args):
             r = self.runtime_info[k]
-            r.set_value(torch_ops.OpRunValue(None if v is None else v.tensor))
+            r.set_value(
+                torch_ops.OpRunValue(None) if v is None else v.__class__(v.tensor_or_sequence)
+            )
 
         # node execution
         for it, kernel in enumerate(self.kernels):
diff --git a/onnx_diagnostic/reference/torch_ops/__init__.py b/onnx_diagnostic/reference/torch_ops/__init__.py
@@ -1,4 +1,4 @@
-from ._op_run import OpRun, OpRunFunction, OpRunValue
+from ._op_run import OpRun, OpRunFunction, OpRunValue, OpRunValueSequence
 from .access_ops import Gather_1, Slice_13
 from .binary_ops import (
     And_1,
@@ -15,11 +15,12 @@
     Pow_12,
     Sub_1,
 )
-from .control_flow import If_1
+from .control_flow import If_1, Loop_16
 from .generator_ops import Range_11
 from .nn_ops import LayerNormalization_17, Softmax_13, Tanh_6
 from .other_ops import Cast_6, CastLike_15, Concat_1, Transpose_1, Trilu_14, Where_9
 from .reduce_ops import ReduceMax_18, ReduceMean_18, ReduceMin_18, ReduceSum_18
+from .sequence_ops import ConcatFromSequence_11, SequenceEmpty_11, SequenceInsert_11
 from .shape_ops import (
     ConstantOfShape_9,
     Expand_8,
diff --git a/onnx_diagnostic/reference/torch_ops/_op_run.py b/onnx_diagnostic/reference/torch_ops/_op_run.py
@@ -1,4 +1,4 @@
-from typing import Any, Optional, Union, Tuple
+from typing import Any, List, Optional, Union, Tuple
 import onnx
 import torch
 from ...helpers import string_type
@@ -18,6 +18,7 @@ class OpRunValue:
     __slots__ = ("cached", "is_constant", "tensor")
 
     def __init__(self, tensor, is_constant: bool = False, may_cpu: bool = False):
+        assert isinstance(tensor, torch.Tensor), f"Unexpected type {type(tensor)}"
         self.tensor = (
             tensor.cpu()
             if may_cpu
@@ -30,6 +31,11 @@ def __init__(self, tensor, is_constant: bool = False, may_cpu: bool = False):
         self.is_constant = is_constant
         self.cached: Optional[Tuple[int, ...]] = None
 
+    @classmethod
+    def is_sequence(cls) -> bool:
+        "Tells if it is sequence."
+        return False
+
     def to(self, to: Any) -> "OpRunValue":
         "Changes the device."
         return OpRunValue(self.tensor.to(to))
@@ -50,6 +56,11 @@ def __repr__(self) -> str:
             )
         return f"{self.__class__.__name__}({string_type(self.tensor, with_shape=True)})"
 
+    @property
+    def tensor_or_sequence(self) -> Union[torch.Tensor, List[torch.Tensor]]:
+        "Returns either a tensor or a sequence."
+        return self.tensor
+
     @property
     def shape(self):
         "shape (torch shape)"
@@ -86,6 +97,41 @@ def as_tuple_int(self) -> Tuple[int, ...]:
         return self._tensor_as_tuple_int()
 
 
+class OpRunValueSequence(OpRunValue):
+    """Defines a sequence."""
+
+    __slots__ = ("cached", "is_constant", "sequence", "tensor")
+
+    def __init__(
+        self, sequence: Optional[List[torch.Tensor]] = None, dtype: torch.dtype = torch.float32
+    ):
+        super().__init__(torch.empty((), dtype=dtype), False, False)
+        self.sequence = sequence or []
+
+    @property
+    def tensor_or_sequence(self) -> Union[torch.Tensor, List[torch.Tensor]]:
+        "Returns either a tensor or a sequence."
+        return self.sequence
+
+    @classmethod
+    def is_sequence(cls) -> bool:
+        "Tells if it is sequence."
+        return True
+
+    def insert_at(
+        self, tensor: OpRunValue, position: Optional[OpRunValue] = None
+    ) -> "OpRunValueSequence":
+        "Inserts a value at a given position."
+        new_seq = OpRunValueSequence()
+        seq = self.sequence.copy()
+        new_seq.sequence = seq
+        if position is None:
+            seq.append(tensor)
+        else:
+            seq.insert(int(position.tensor.item()), tensor)
+        return new_seq
+
+
 class OpRun:
     """
     Main class. Every kernel should inherit from it.
diff --git a/onnx_diagnostic/reference/torch_ops/control_flow.py b/onnx_diagnostic/reference/torch_ops/control_flow.py
diff --git a/onnx_diagnostic/reference/torch_ops/sequence_ops.py b/onnx_diagnostic/reference/torch_ops/sequence_ops.py
diff --git a/onnx_diagnostic/torch_onnx/runtime_info.py b/onnx_diagnostic/torch_onnx/runtime_info.py
