more checks

Author: xadupre

onnx_diagnostic/reference/ort_evaluator.py

@@ -81,6 +81,10 @@ def to(self, tensor_like) -> "OnnxList":
             ]
         )
 
+    def clone(self) -> "OnnxList":
+        "Clone (torch)."
+        return [t.clone() for t in self]
+
 
 class OnnxruntimeEvaluator:
     """
@@ -708,6 +712,10 @@ def _run(self, node: NodeProto, inputs: List[Any], results: Dict[str, Any]) -> L
 
         outputs = list(sess.run(None, feeds))
         assert isinstance(outputs, list), f"Unexpected type for outputs {type(outputs)}"
+        assert not any(type(v) is list for v in outputs), (
+            f"One output type is a list, this should not be allowed, "
+            f"node.op_type={node.op_type}, feeds={string_type(feeds,with_shape_type=True)}"
+        )
         return outputs
 
     def _run_if(
@@ -783,6 +791,11 @@ def _run_scan_or_loop(
         self, node: NodeProto, inputs: List[Any], results: Dict[str, Any]
     ) -> List[Any]:
         """Runs a node Scan."""
+        assert not any(type(i) is list for i in inputs), (
+            f"One input is a list but it should an OnnxList, "
+            f"node.op_type={node.op_type!r}, node.input={node.input}, "
+            f"inputs={string_type(inputs, with_shape=True)}"
+        )
         feeds = dict(zip(node.input, inputs))
         feeds.update(results)
         name = "body"

onnx_diagnostic/torch_onnx/sbs.py

@@ -9,6 +9,7 @@
 from ..helpers.onnx_helper import pretty_onnx
 from ..helpers.torch_helper import to_numpy, from_numpy, to_tensor, torch_dtype_to_onnx_dtype
 from ..helpers.torch_fx_graph_helper import prepare_args_kwargs, run_fx_node
+from ..reference.ort_evaluator import OnnxList
 from .sbs_dataclasses import (
     ReplayConfiguration,
     RunAlignedRecord,
@@ -26,11 +27,11 @@ def _check_tensor_(use_tensor, name, obj, flip_type=False):
             if isinstance(obj, torch.Tensor):
                 obj = to_numpy(obj)
 
-    assert not use_tensor or isinstance(obj, torch.Tensor), (
+    assert not use_tensor or isinstance(obj, (torch.Tensor, OnnxList)), (
         f"Unexpected type {type(obj)} for {name!r}. "
         f"use_tensor is True so torch.Tensor is expected."
     )
-    assert use_tensor or isinstance(obj, np.ndarray), (
+    assert use_tensor or isinstance(obj, (np.ndarray, OnnxList)), (
         f"Unexpected type {type(obj)} for {name!r}. "
         f"use_tensor is False so np.array is expected."
     )