fix sbs

Author: xadupre

.github/workflows/check-urls.yml

@@ -42,6 +42,6 @@ jobs:
         print_all: false
         timeout: 2
         retry_count# : 2
-        exclude_urls: https://hal.archives-,ouvertes.fr/hal-00990252/document,http://badge.fury.io/py/onnx-diagnostic,https://azure.microsoft.com/en-us/products/devops/pipelines,https://github.com/pytorch/pytorch/pull/83137#issuecomment-1211320670,https://github.com/NVIDIA/TransformerEngine.git@6a9edc38bf9b941b7d369af5103fa8fe0b121d61,https://medium.com/@msouza.os/llm-from-scratch-with-pytorch-9f21808c6319,https://github.com/pytorch/pytorch/blob/main/torch/fx/experimental/symbolic_shapes.py#L5965,https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-24-04.html,https://badge.fury.io/py/onnx-diagnostic.svg,https://github.com/huggingface/transformers/pull/36311,https://www.linux.org/
-        exclude_patterns: https://www.data.gouv.fr/fr/datasets/r/e3d83ab3-dc52-4c99-abaf-8a38050cc68c,https://dev.azure.com/,https://azure.microsoft.com/en-us/products/devops/pipelines,https://github.com/pytorch/pytorch/pull/83137#issuecomment-1211320670,https://github.com/NVIDIA/TransformerEngine.git@6a9edc38bf9b941b7d369af5103fa8fe0b121d61,https://github.com/pytorch/pytorch/blob/main/torch/,https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-24-04.html,https://badge.fury.io/py/onnx-diagnostic.svg,https://github.com/,https://codecov.io/,https://huggingface.co/,https://www.linux.org/
+        exclude_urls: https://hal.archives-,ouvertes.fr/hal-00990252/document,http://badge.fury.io/py/onnx-diagnostic,https://azure.microsoft.com/en-us/products/devops/pipelines,https://github.com/pytorch/pytorch/pull/83137#issuecomment-1211320670,https://github.com/NVIDIA/TransformerEngine.git@6a9edc38bf9b941b7d369af5103fa8fe0b121d61,https://medium.com/@msouza.os/llm-from-scratch-with-pytorch-9f21808c6319,https://github.com/pytorch/pytorch/blob/main/torch/fx/experimental/symbolic_shapes.py#L5965,https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-24-04.html,https://badge.fury.io/py/onnx-diagnostic.svg,https://github.com/huggingface/transformers/pull/36311,https://www.linux.org/,https://docs.scipy.org/doc/scipy/
+        exclude_patterns: https://www.data.gouv.fr/fr/datasets/r/e3d83ab3-dc52-4c99-abaf-8a38050cc68c,https://dev.azure.com/,https://azure.microsoft.com/en-us/products/devops/pipelines,https://github.com/pytorch/pytorch/pull/83137#issuecomment-1211320670,https://github.com/NVIDIA/TransformerEngine.git@6a9edc38bf9b941b7d369af5103fa8fe0b121d61,https://github.com/pytorch/pytorch/blob/main/torch/,https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-24-04.html,https://badge.fury.io/py/onnx-diagnostic.svg,https://github.com/,https://codecov.io/,https://huggingface.co/,https://www.linux.org/,https://docs.scipy.org/doc/scipy/
         # force_pass : true

onnx_diagnostic/reference/ort_evaluator.py

@@ -473,7 +473,7 @@ def enumerate_nodes(self, nodes: List[NodeProto]) -> Iterator[NodeProto]:
             yield node
 
     @classmethod
-    def _get_hidden_inputs(self, graph: GraphProto) -> Set[str]:
+    def _get_hidden_inputs(cls, graph: GraphProto) -> Set[str]:
         """
         Returns the hidden inputs (inputs coming from an upper context)
         used by a subgraph.
@@ -490,21 +490,21 @@ def _get_hidden_inputs(self, graph: GraphProto) -> Set[str]:
                     hidden.add(i)
             for att in node.attribute:
                 if att.type == AttributeProto.GRAPH and att.g:
-                    hid = self._get_hidden_inputs(att.g)
+                    hid = cls._get_hidden_inputs(att.g)
                     less = set(h for h in hid if h not in memo)
                     hidden |= less
             memo |= set(node.output)
         return hidden
 
     @classmethod
-    def _get_hidden_node_inputs(self, node: NodeProto) -> Set[str]:
+    def _get_hidden_node_inputs(cls, node: NodeProto) -> Set[str]:
         """Calls multiple _get_hidden_inputs on every attribute."""
         if node.op_type not in {"Loop", "Scan", "If"}:
             return set()
         hidden = set()
         for att in node.attribute:
             if att.type == AttributeProto.GRAPH:
-                hidden |= self._get_hidden_inputs(att.g)
+                hidden |= cls._get_hidden_inputs(att.g)
         return hidden - (hidden & set(node.input))
 
     def _get_sess(

onnx_diagnostic/torch_onnx/sbs.py

@@ -9,7 +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 ..reference.ort_evaluator import OnnxList, OnnxruntimeEvaluator
 from .sbs_dataclasses import (
     ReplayConfiguration,
     RunAlignedRecord,
@@ -176,10 +176,12 @@ def _loop_onnx_node(
 
     ref = run_cls(node, **run_cls_kwargs)
     # We need to clone because the runtime maybe using dlpack to create OrtValue
+    hidden_inputs = OnnxruntimeEvaluator._get_hidden_node_inputs(node)
+    all_inputs = [*node.input, *hidden_inputs] if hidden_inputs else node.input
     feeds = (
-        {k: onnx_results[k].clone() for k in node.input if k}
+        {k: onnx_results[k].clone() for k in all_inputs if k}
         if use_tensor
-        else {k: onnx_results[k].copy() for k in node.input if k}
+        else {k: onnx_results[k].copy() for k in all_inputs if k}
     )
     assert "" not in feeds, f"Unexpected feeds={string_type(feeds, **str_kws)}"
     if verbose > 1: