onnx-diagnostic 0.8.4__py3-none-any.whl → 0.8.6__py3-none-any.whl

onnx_diagnostic/torch_onnx/compare.py

@@ -0,0 +1,357 @@
+import enum
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Tuple, Union
+import numpy as np
+import onnx
+from ..helpers.onnx_helper import onnx_dtype_name
+
+
+_NOT_SO_FAR_OPS = [
+    {"MatMul", "Gemm", "FusedMatMul"},
+    {"Conv", "FusedConv"},
+    {"MaxPool"},
+]
+
+
+def _sum_sets(sets):
+    t = set()
+    for s in sets:
+        t |= s
+    return t
+
+
+_ALL_NOT_SO_FAR_OPS = _sum_sets(_NOT_SO_FAR_OPS)
+
+
+def _align(res: str, limit: int) -> str:
+    if len(res) == limit:
+        return res
+    if len(res) > limit:
+        return res[:limit]
+    return res + " " * (limit - len(res))
+
+
+class ObsType(enum.IntEnum):
+    """Observation kind."""
+
+    RESULT = 1
+    INITIALIZER = 2
+    SPARSE_INITIALIZER = 4
+    INPUT = 8
+    OUTPUT = 16
+    NODE = 32
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}.{self._name_}"
+
+
+@dataclass
+class ObsCompare:
+    """
+    The description of an observation, a node, an input, an output, an initializer.
+
+    :param position: index of this observation in the original model
+    :param kind: node type, see :class:`ObsType`
+    :param name_or_outputs: name of an initializer or the outputs of a node
+    :param itype: onnx type
+    :param index: index of an input or output
+    :param shape: shape
+    :param op_type: node op_type
+    :param comment: comment, unused
+    """
+
+    position: int
+    kind: ObsType
+    name_or_outputs: Tuple[str]
+    itype: int = 0
+    index: int = 0
+    shape: Optional[Tuple[Tuple[Union[int, str], ...]]] = None
+    op_type: str = ""
+    comment: str = ""
+
+    def __str__(self) -> str:
+        "usual"
+        els = [
+            _align(f"{self.position:04d}", 4),
+            _align(self.kind._name_, 6),
+            _align(onnx_dtype_name(self.itype) if self.itype else "?", 8),
+            _align("?" if self.shape is None else "x".join(map(str, self.shape)), 18),
+            _align(self.op_type or "", 15),
+            _align(", ".join(self.name_or_outputs), 35),
+        ]
+        return " ".join(els)
+
+    @classmethod
+    def to_str(cls, obs: Optional["ObsCompare"]) -> str:
+        assert not obs or isinstance(obs, ObsCompare), f"unexpected type {type(obs)}"
+        if obs:
+            return str(obs)
+        return " " * (4 + 6 + 8 + 18 + 15 + 35 + 5)
+
+    def distance(self, obs: "ObsCompare") -> float:
+        """Computes a cost between two observations."""
+        if self.kind != obs.kind:
+            return 1e6
+        d: float = 0
+        if self.itype != obs.itype:
+            d += 1e5
+        if self.kind == ObsType.NODE:
+            cost = 9997
+            d = 0
+            if self.op_type != obs.op_type:
+                if self.op_type in _ALL_NOT_SO_FAR_OPS or obs.op_type in _ALL_NOT_SO_FAR_OPS:
+                    d += 1e2
+                    for aset in _NOT_SO_FAR_OPS:
+                        if self.op_type in aset and obs.op_type in aset:
+                            cost = 97
+                        elif self.op_type in aset or obs.op_type in aset:
+                            d += 5e4
+                else:
+                    d += 9e2
+            if len(self.name_or_outputs) == 1 and len(obs.name_or_outputs) == 1:
+                if self.name_or_outputs[0] != obs.name_or_outputs[0]:
+                    n1 = self.name_or_outputs[0]
+                    n2 = obs.name_or_outputs[0]
+                    n1 = n1.replace("_", "")
+                    n2 = n2.replace("_", "")
+                    if n1 == n2:
+                        d += 1
+                    elif (n1.startswith(("val_", "_onx_")) or "::" in n1 or "--" in n1) and (
+                        n2.startswith(("val_", "_onx_")) or "::" in n2 or "--" in n2
+                    ):
+                        # These are name given the exporter
+                        # and not inspired from the model itself.
+                        d += cost / 100
+                    else:
+                        d += cost
+            else:
+                a = set(self.name_or_outputs) & set(obs.name_or_outputs)
+                b = set(self.name_or_outputs) | set(obs.name_or_outputs)
+                d += cost * (len(b) - len(a))
+            return d
+        if self.kind == ObsType.INPUT:
+            return (
+                999.7
+                if self.itype != obs.itype
+                or self.shape != obs.shape
+                or self.index != obs.index
+                else 0
+            )
+        if self.kind == ObsType.INITIALIZER or self.kind == ObsType.SPARSE_INITIALIZER:
+            return 1e3 if self.itype != obs.itype or self.shape != obs.shape else 0
+        if self.kind == ObsType.OUTPUT:
+            return (
+                999.1
+                if self.itype != obs.itype
+                or self.shape != obs.shape
+                or self.index != obs.index
+                else 0
+            )
+        return 1e8
+
+    @classmethod
+    def obs_sequence_from_model(
+        cls,
+        model: Union[onnx.ModelProto, onnx.GraphProto],
+    ) -> List["ObsCompare"]:
+        """
+        Creates a sequence of observations bases on a model.
+
+        :param model: model
+        :return: sequence of observations
+        """
+        graph = model if isinstance(model, onnx.GraphProto) else model.graph
+
+        shapes = {}
+        types = {}
+        for info in [*graph.value_info, *graph.input, *graph.output]:
+            if info.type.tensor_type:
+                t = info.type.tensor_type
+                shapes[info.name] = tuple((d.dim_param or d.dim_value) for d in t.shape.dim)
+                types[info.name] = t.elem_type
+
+        seq: List[ObsCompare] = []
+        for init in graph.initializer:
+            obs = ObsCompare(
+                position=len(seq),
+                kind=ObsType.INITIALIZER,
+                itype=init.data_type,
+                shape=tuple(init.dims),
+                name_or_outputs=(init.name,),
+            )
+            seq.append(obs)
+        for i, inp in enumerate(graph.input):
+            obs = ObsCompare(
+                position=len(seq),
+                kind=ObsType.INPUT,
+                itype=inp.type.tensor_type.elem_type,
+                index=i,
+                shape=tuple(
+                    (d.dim_param or d.dim_value) for d in inp.type.tensor_type.shape.dim
+                ),
+                name_or_outputs=(inp.name,),
+            )
+            seq.append(obs)
+        for node in graph.node:
+            obs = ObsCompare(
+                position=len(seq),
+                kind=ObsType.NODE,
+                itype=types.get(node.output[0], 0),
+                index=i,
+                shape=shapes.get(node.output[0], None),
+                name_or_outputs=tuple(node.output),
+                op_type=node.op_type,
+            )
+            seq.append(obs)
+        for i, inp in enumerate(graph.output):
+            obs = ObsCompare(
+                position=len(seq),
+                kind=ObsType.OUTPUT,
+                itype=inp.type.tensor_type.elem_type,
+                index=i,
+                shape=tuple(
+                    (d.dim_param or d.dim_value) for d in inp.type.tensor_type.shape.dim
+                ),
+                name_or_outputs=(inp.name,),
+            )
+            seq.append(obs)
+        return seq
+
+
+@dataclass
+class ObsComparePair:
+    """
+    Defines a pair of comparison objects
+
+    :param side1: object from first side
+    :param side2: object from first side
+    :param distance: distance
+    """
+
+    side1: Optional[ObsCompare]
+    side2: Optional[ObsCompare]
+    distance: float
+
+    def __str__(self) -> str:
+        "nice display"
+        return (
+            f"{self.distance:.4e} | "
+            f"{ObsCompare.to_str(self.side1)} | {ObsCompare.to_str(self.side2)}"
+        )
+
+    @classmethod
+    def to_str(cls, seq: List["ObsComparePair"]) -> str:
+        """Displays every pair in text."""
+        return "\n".join([f"{str(pair)}" for pair in seq])
+
+    @classmethod
+    def distance_sequence(cls, s1: List["ObsCompare"], s2: List["ObsCompare"]) -> Tuple[
+        float,
+        List[Tuple[int, int]],
+        List["ObsComparePair"],
+    ]:
+        """
+        Computes the distance between two sequences of results.
+
+        :param s1: first sequence
+        :param s2: second sequence
+        :return: distance and alignment
+
+        An example:
+
+        .. runpython::
+            :showcode:
+
+            import torch
+            from onnx_diagnostic.export.api import to_onnx
+            from onnx_diagnostic.torch_onnx.compare import ObsComparePair, ObsCompare
+
+
+            class Model(torch.nn.Module):
+                def __init__(self):
+                    super().__init__()
+                    self.conv1 = torch.nn.Conv2d(3, 16, 5)
+                    self.fc1 = torch.nn.Linear(144, 64)
+                    self.fc2 = torch.nn.Linear(64, 128)
+                    self.fc3 = torch.nn.Linear(128, 10)
+
+                def forward(self, x):
+                    x = torch.nn.functional.max_pool2d(
+                        torch.nn.functional.relu(self.conv1(x)),
+                        (4, 4),
+                    )
+                    # x = F.max_pool2d(F.relu(self.conv2(x)), 2)
+                    x = torch.flatten(x, 1)
+                    x = torch.nn.functional.relu(self.fc1(x))
+                    x = torch.nn.functional.relu(self.fc2(x))
+                    y = self.fc3(x)
+                    return y
+
+
+            model = Model()
+            x = torch.randn((2, 3, 16, 17), dtype=torch.float32)
+            dynamic_shapes = ({0: "batch", 3: "dim"},)
+            onnx_optimized = to_onnx(
+                model, (x,), dynamic_shapes=dynamic_shapes, exporter="custom", optimize=True
+            ).model_proto
+            onnx_not_optimized = to_onnx(
+                model, (x,), dynamic_shapes=dynamic_shapes, exporter="custom", optimize=False
+            ).model_proto
+            seq1 = ObsCompare.obs_sequence_from_model(onnx_not_optimized)
+            seq2 = ObsCompare.obs_sequence_from_model(onnx_optimized)
+            _dist, _path, pair_cmp = ObsComparePair.distance_sequence(seq1, seq2)
+            text = ObsComparePair.to_str(pair_cmp)
+            print(text)
+        """
+        delay = max(50, abs(len(s2) - len(s1)) + 1)
+        distance: Dict[Tuple[int, int], Union[int, float]] = {(-1, -1): 0}
+        predecessor: Dict[Tuple[int, int], Optional[Tuple[int, int]]] = {(-1, -1): None}
+        insert_cost = 1e3
+        for i in range(len(s1)):
+            for j in range(max(0, i - delay), min(len(s2), i + delay)):
+                best = distance.get((i, j), 1e100)
+                pred = None
+                ki, kj = i - 1, j - 1
+                if (ki, kj) in distance:
+                    d = distance[ki, kj] + s1[i].distance(s2[j])
+                    if d < best:
+                        best = d
+                        pred = (ki, kj)
+                ki, kj = i - 1, j
+                if (ki, kj) in distance:
+                    d = distance[ki, kj] + insert_cost + 1
+                    if d < best:
+                        best = d
+                        pred = (ki, kj)
+                ki, kj = i, j - 1
+                if (ki, kj) in distance:
+                    d = distance[ki, kj] + insert_cost + 0.1
+                    if d < best:
+                        best = d
+                        pred = (ki, kj)
+                distance[i, j] = best
+                predecessor[i, j] = pred
+
+        # reverse
+        way = []
+        last: Optional[Tuple[int, int]] = len(s1) - 1, len(s2) - 1
+        while last is not None:
+            way.append(last)
+            last = predecessor[last]
+        indices = list(reversed(way))[1:]
+        obs_path: List[ObsComparePair] = []
+        last = -1, -1
+        for i, j in indices:
+            di = i - last[0]
+            dj = j - last[1]
+            cost = distance.get((i, j), np.nan)
+            if di == dj == 1:
+                obs_path.append(ObsComparePair(s1[i], s2[j], distance=cost))
+            elif di == 0:
+                obs_path.append(ObsComparePair(None, s2[j], distance=cost))
+            elif dj == 0:
+                obs_path.append(ObsComparePair(s1[i], None, distance=cost))
+            else:
+                raise RuntimeError(f"issue with di={di}, dj={dj}")
+            last = i, j
+        return distance[len(s1) - 1, len(s2) - 1], indices, obs_path

{onnx_diagnostic-0.8.4.dist-info → onnx_diagnostic-0.8.6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: onnx-diagnostic
-Version: 0.8.4
+Version: 0.8.6
 Summary: Tools to help converting pytorch models into ONNX.
 Home-page: https://github.com/sdpython/onnx-diagnostic
 Author: Xavier Dupré

{onnx_diagnostic-0.8.4.dist-info → onnx_diagnostic-0.8.6.dist-info}/RECORD

@@ -1,16 +1,18 @@
-onnx_diagnostic/__init__.py,sha256=BuvSD4fYmz8ZVadwiG4GCHeb46p5sNX7-_GM16OtKW0,173
+onnx_diagnostic/__init__.py,sha256=YQit5D2idhb9-wNQZzvWLT_qwRrKWBoTqMpNlBaWsGw,173
 onnx_diagnostic/__main__.py,sha256=YmyV_Aq_ianDlHyKLHMa6h8YK3ZmFPpLVHLKjM91aCk,79
-onnx_diagnostic/_command_lines_parser.py,sha256=Xkh_7fIDbT5ghTpLqlVhx4cIxAUXqv6zvPdnN3aCdOY,52254
+onnx_diagnostic/_command_lines_parser.py,sha256=AWT6XrphbR0C0w9J846jPcRWkoUtnSSAX7gdR-JavQ4,54258
 onnx_diagnostic/api.py,sha256=BhCl_yCd78N7TlVtPOHjeYv1QBEy39TjZ647rcHqLh0,345
 onnx_diagnostic/doc.py,sha256=t3RELgfooYnVMAi0JSpggWkQEgUsREz8NmRvn0TnLI8,2829
-onnx_diagnostic/ext_test_case.py,sha256=rVZWqFEfnvwnsD3wF4jeDblh5uj5ckZ8C6DZQ0RGb_E,49599
+onnx_diagnostic/ext_test_case.py,sha256=KxRC6s9107hYvNgU9x2B85rj8_EhAtymPIlMpmkUNu8,50154
+onnx_diagnostic/ci_models/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+onnx_diagnostic/ci_models/ci_helpers.py,sha256=6CKQ4dVHBHeF6rN_Q3Y_0ZFeLYbYQbGQO3YyW3PQyAc,15341
+onnx_diagnostic/ci_models/export_qwen25_vl.py,sha256=pyxmMIps9aDNkXzDAFrG8Q9DDOsEyKRHjoVvAggjFdU,20050
 onnx_diagnostic/export/__init__.py,sha256=yEIoWiOeTwBsDhyYt2fTKuhtA0Ya1J9u9ZzMTOTWaWs,101
-onnx_diagnostic/export/api.py,sha256=c3ZASq2upUAoiQ4aymm8vOAEySN_Yk6l0o1hWf6Ailo,10065
-onnx_diagnostic/export/control_flow.py,sha256=zU5n_QYhNcBllyMsl1_i6ohZt2CshqG2MokJghrvA60,7751
-onnx_diagnostic/export/control_flow_onnx.py,sha256=sODOD4v7EJj6LWhrfcdCW68r9nYKsRM4SRnqDw4TrSI,18049
-onnx_diagnostic/export/control_flow_research.py,sha256=RuYz9_eM42Bk6TKSiPV6dS68LIMZu-6WBCFCKoSvjrk,5422
+onnx_diagnostic/export/api.py,sha256=BX4c99gMlRYsBWk3P15FMRogArxjP4dXYXP5gILjgIk,10626
+onnx_diagnostic/export/cf_simple_loop_for.py,sha256=0I1tRAwhmmqA-6Qaq8AiUL0Ci-HODuRAVcI9azNcxAQ,13345
+onnx_diagnostic/export/control_flow_onnx.py,sha256=izGlctqQANrHzSxPMbT7hoauNbnIBdx6hb8ry7HtVmM,18263
 onnx_diagnostic/export/dynamic_shapes.py,sha256=M2hlpHSTbkzZwGKAbrpQXng5HQrwjF5Z6wGGxEgnp74,42061
-onnx_diagnostic/export/onnx_plug.py,sha256=WqqdTBk2pV26AplNQysIdhR9y3ZFdQ-5KXu5ogTNcgI,21053
+onnx_diagnostic/export/onnx_plug.py,sha256=U13fL0BjnhMzcDGxaAOqM4TQte5Z4zKDg4ESS0iktjM,22704
 onnx_diagnostic/export/shape_helper.py,sha256=m628y0oRCQbeZkeh8JDHIfWMsSjoJoeX-IPiPGDHT-w,11273
 onnx_diagnostic/export/validate.py,sha256=_PGUql2DJhIgGKo0WjTGUc5AgsZUx8fEs00MePy-w98,6043
 onnx_diagnostic/helpers/__init__.py,sha256=GJ2GT7cgnlIveVUwMZhuvUwidbTJaKv8CsSIOpZDsJg,83
@@ -23,7 +25,7 @@ onnx_diagnostic/helpers/doc_helper.py,sha256=pl5MZd3_FaE8BqQnqoBuSBxoNCFcd2OJd3e
 onnx_diagnostic/helpers/dot_helper.py,sha256=hwgTJsbsUv0qq7euyPDnc1NsBZDGOwv32JXSZxIHJkE,8118
 onnx_diagnostic/helpers/fake_tensor_helper.py,sha256=J7wnK3WTuVKnYiMzLVTAPkdJr3hQfIfMC9ZlOu7oGmI,11024
 onnx_diagnostic/helpers/graph_helper.py,sha256=hevQT5a7_QuriVPQcbT5qe18n99Doyl5h3-qshx1-uk,14093
-onnx_diagnostic/helpers/helper.py,sha256=f5w53QR0IO1-zqAMacgeGpZNA8uAo0c2k_ZYXP_BRhE,65840
+onnx_diagnostic/helpers/helper.py,sha256=x8EYQmgrz_G5QS_IsbeFIoDcN_sUs-CslJMHseBj1Fw,65482
 onnx_diagnostic/helpers/log_helper.py,sha256=0lJiTF87lliI-LmgpUH_V2N8NuzJ0LryH0mSYpkRaL8,93272
 onnx_diagnostic/helpers/memory_peak.py,sha256=M3m4_thWFIwP5HytbJYEqaijXIv5v5BW_vlcJowIYI4,6434
 onnx_diagnostic/helpers/mini_onnx_builder.py,sha256=jR2lkRZEQ0N30H0FqeBwaxJd_w_6kyxFagrnulqFjhE,23883
@@ -100,7 +102,7 @@ onnx_diagnostic/tasks/zero_shot_image_classification.py,sha256=jJCMWuOqGv5ahCfjr
 onnx_diagnostic/tasks/data/__init__.py,sha256=uJoemrWgEjI6oA-tMX7r3__x-b3siPmkgqaY7bgIles,401
 onnx_diagnostic/tasks/data/dummies_imagetext2text_generation_gemma3.onnx,sha256=UbtvmWMqcZOKJ-I-HXWI1A6YR6QDaFS5u_yXm5C3ZBw,10299
 onnx_diagnostic/torch_export_patches/__init__.py,sha256=0SaZedwznm1hQUCvXZsGZORV5vby954wEExr5faepGg,720
-onnx_diagnostic/torch_export_patches/onnx_export_errors.py,sha256=WPb8Ku643UIV8kDyt9JUpaJBIVXth9UbteCNctd_yis,41863
+onnx_diagnostic/torch_export_patches/onnx_export_errors.py,sha256=OpZHNWiA0iU-6WCFZcVCj06_MopYiZQ6c6CbAuSQ8Ms,42357
 onnx_diagnostic/torch_export_patches/onnx_export_serialization.py,sha256=0HdubI06EGpxOICqDWZoVmZkVO9gAaFADEmby197EyM,11935
 onnx_diagnostic/torch_export_patches/patch_details.py,sha256=MSraVo5ngBhihi8ssPMXSY9B4fJ17J-GAADaw3dT-rc,11794
 onnx_diagnostic/torch_export_patches/patch_expressions.py,sha256=vr4tt61cbDnaaaduzMj4UBZ8OUtr6GfDpIWwOYqjWzs,3213
@@ -117,15 +119,15 @@ onnx_diagnostic/torch_export_patches/patches/_patch_transformers_dynamic_cache.p
 onnx_diagnostic/torch_export_patches/patches/_patch_transformers_gemma3.py,sha256=nVgYQk0xXpHiictN1wOHVMN2lTH9b0vfIJ4ie-uKopg,1999
 onnx_diagnostic/torch_export_patches/patches/_patch_transformers_generation_mixin.py,sha256=VIZsVHgR8NmAcBQalPl5I6ZzNgcBxjGb6ars31m9gRg,21936
 onnx_diagnostic/torch_export_patches/patches/_patch_transformers_idefics.py,sha256=kTjuTRsfkGGGhspJnMxAMQSchZgGC_IruJzpHh_FmI8,6348
-onnx_diagnostic/torch_export_patches/patches/_patch_transformers_masking_utils.py,sha256=R4YwnN9ktxjjImiJtLRxiKtKLr9LuFlwkPXkTJ6BTIo,6895
+onnx_diagnostic/torch_export_patches/patches/_patch_transformers_masking_utils.py,sha256=HE3fovyvMiYe9EPz1UjdD9AWopX3H188SMwPb8w5mzM,7111
 onnx_diagnostic/torch_export_patches/patches/_patch_transformers_qwen2.py,sha256=OxYdlLrwtd_KGHt3E17poduxvWFg-CfGS57-yN1i6gI,3827
-onnx_diagnostic/torch_export_patches/patches/_patch_transformers_qwen2_5.py,sha256=icjSrI3LFrShtV_AYQ8F2qiMFHZ74Qg5I2c-V23uEgg,31601
+onnx_diagnostic/torch_export_patches/patches/_patch_transformers_qwen2_5.py,sha256=yALbXWi3ysJ6nzQD-rxTdxdNJiBsTbYEBIj4TdksDOA,34598
 onnx_diagnostic/torch_export_patches/patches/_patch_transformers_qwen3.py,sha256=cND9Iqo1aKdlX-BXGr9Qlq_Y4EW1L5VWSwZfqYTVazU,4888
 onnx_diagnostic/torch_export_patches/patches/_patch_transformers_rotary_embedding.py,sha256=4bJ_z2gizZQla_fcCVt0dmuhzO9Vu-D7CCMWdxMlrKM,16893
 onnx_diagnostic/torch_export_patches/patches/_patch_transformers_sam_mask_decoder.py,sha256=-6TuBm3sLAFEGuW3vRfOTtE5uP6aINFfu7xMnl27Dws,5703
 onnx_diagnostic/torch_export_patches/patches/patch_helper.py,sha256=kK_CGW643iVXxa-m6pttDBS7HTyMQaPypza7iqIInn4,721
 onnx_diagnostic/torch_export_patches/patches/patch_torch.py,sha256=FfES0WWiWxmuQbGTlQ7IJS0YBG7km3IQbnMYwk_lPPU,44667
-onnx_diagnostic/torch_export_patches/patches/patch_transformers.py,sha256=VAfZj0xu3D4CG71SWv-9sYPUK4ZQTSz2-x4qxP4DxGE,3079
+onnx_diagnostic/torch_export_patches/patches/patch_transformers.py,sha256=Mvq8q1Lz3l3GyCD6j8WQjbrPk_V2dnc4iKm3cC_o1OA,3112
 onnx_diagnostic/torch_export_patches/serialization/__init__.py,sha256=BHLdRPtNAtNPAS-bPKEj3-foGSPvwAbZXrHzGGPDLEw,1876
 onnx_diagnostic/torch_export_patches/serialization/diffusers_impl.py,sha256=drq3EH_yjcSuIWYsVeUWm8Cx6YCZFU6bP_1PLtPfY5I,945
 onnx_diagnostic/torch_export_patches/serialization/transformers_impl.py,sha256=sIHFvUQoMK8ytXQYB-k7OL62z8A3f5uDaq-S5R5uN-M,10034
@@ -136,18 +138,19 @@ onnx_diagnostic/torch_models/validate.py,sha256=fnbTl5v1n5nM2MpmCgCMaWa6c7DGpb5m
 onnx_diagnostic/torch_models/hghub/__init__.py,sha256=vi1Q7YHdddj1soiBN42MSvJdFqe2_KUoWafHISjwOu8,58
 onnx_diagnostic/torch_models/hghub/hub_api.py,sha256=rFbiPNLET-KdBpnv-p0nKgwHX6d7C_Z0s9zZ86_92kQ,14307
 onnx_diagnostic/torch_models/hghub/hub_data.py,sha256=8V_pAgACPLPsLRYUododg7MSL6str-T3tBEGY4OaeYQ,8724
-onnx_diagnostic/torch_models/hghub/hub_data_cached_configs.py,sha256=GimzkI8W3guATkDx7RQ-w2xNGVaFDVegfTnnmNxf4iE,292068
+onnx_diagnostic/torch_models/hghub/hub_data_cached_configs.py,sha256=Dxa13rsnTQ8eH_BcQvbY2bp1AYFtzuFrJ-J_urrSmeQ,292694
 onnx_diagnostic/torch_models/hghub/model_inputs.py,sha256=tCGqigRyY1omxm2rczRUvCTsweZGbF1MccWI3MmCH20,17423
 onnx_diagnostic/torch_models/hghub/model_specific.py,sha256=j50Nu7wddJMoqmD4QzMbNdFDUUgUmSBKRzPDH55TlUQ,2498
 onnx_diagnostic/torch_models/untrained/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 onnx_diagnostic/torch_models/untrained/llm_phi2.py,sha256=y_akbdApi136qHcEQgykwIAYVw0Yfi0lbjb3DNuafaU,3948
 onnx_diagnostic/torch_models/untrained/llm_tiny_llm.py,sha256=QXw_Bs2SzfeiQMf-tmtVl83SmVOL4-Um7Qy-f0E48QI,2507
 onnx_diagnostic/torch_onnx/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+onnx_diagnostic/torch_onnx/compare.py,sha256=O0lws4kzn8WAXr8-x-YMPr7oyBC9DtSIs4OfOr4S5-E,12305
 onnx_diagnostic/torch_onnx/runtime_info.py,sha256=u1bD6VXqzBCRmqmbzQtDswaPs1PH_ygr1r-CrcfXpNU,8562
 onnx_diagnostic/torch_onnx/sbs.py,sha256=8okBEIupMgw7TtKc80YFimMtwnY3GchdY05FsA9ooa0,40749
 onnx_diagnostic/torch_onnx/sbs_dataclasses.py,sha256=UctdBjzoPTQG1LS0tZ8A6E9hpoq5HWUYaJLPOPJc9FI,20299
-onnx_diagnostic-0.8.4.dist-info/licenses/LICENSE.txt,sha256=Vv6TXglX6Rc0d-f8aREhayhT-6PMQXEyOmI2NKlUCMc,1045
-onnx_diagnostic-0.8.4.dist-info/METADATA,sha256=8S9bFx2lTef7dTeL_dTCtGj1MalIyoUvs5dMzrMffNg,6734
-onnx_diagnostic-0.8.4.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-onnx_diagnostic-0.8.4.dist-info/top_level.txt,sha256=KwNkXewmcobM3ZT1DJLVWH6ebJzA5qKg7cWqKfpGNT4,16
-onnx_diagnostic-0.8.4.dist-info/RECORD,,
+onnx_diagnostic-0.8.6.dist-info/licenses/LICENSE.txt,sha256=Vv6TXglX6Rc0d-f8aREhayhT-6PMQXEyOmI2NKlUCMc,1045
+onnx_diagnostic-0.8.6.dist-info/METADATA,sha256=9xPlJ9UHYSSIyEMqxN14mqZg31Rq9jqzHaczfqWhu-4,6734
+onnx_diagnostic-0.8.6.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+onnx_diagnostic-0.8.6.dist-info/top_level.txt,sha256=KwNkXewmcobM3ZT1DJLVWH6ebJzA5qKg7cWqKfpGNT4,16
+onnx_diagnostic-0.8.6.dist-info/RECORD,,

onnx_diagnostic/export/control_flow.py

@@ -1,214 +0,0 @@
-import contextlib
-from typing import Any, Callable, List, Optional, Sequence, Tuple, Union
-import torch
-from torch._higher_order_ops.utils import (
-    materialize_as_graph,
-    check_input_alias_and_mutation_return_outputs,
-    # _maybe_reenter_make_fx,
-)
-
-_TEST_EXPORT = False
-
-
-@contextlib.contextmanager
-def enable_code_export_control_flow():
-    """Enables the code meant to be exported."""
-    global _TEST_EXPORT
-    old = _TEST_EXPORT
-    _TEST_EXPORT = True
-    try:
-        yield
-    finally:
-        _TEST_EXPORT = old
-
-
-def is_exporting() -> bool:
-    """
-    Returns :func:`torch.compiler.is_exporting` or
-    :func:`torch.compiler.is_compiling`.
-    Changes ``_TEST_EXPORT`` to make it trigger.
-    """
-    return _TEST_EXPORT or torch.compiler.is_exporting() or torch.compiler.is_compiling()
-
-
-def _loop_for_fn(n_iter, body_fn, reduction_dim, args):
-    """
-    Python implementation of the loop.
-
-    :param n_iter: number of iteration
-    :param body_fn: function implementing the body
-    :param reduction_dim: dimension used to reduce the list produced by the loop
-    :param args: arguments to the loop body
-    :return: results
-    """
-    res = []
-    for i in torch.arange(n_iter, dtype=n_iter.dtype):
-        r = body_fn(i, *args)
-        if isinstance(r, tuple):
-            assert not res or len(r) == len(res[-1]), (
-                f"Unexpected number of results {len(r)} for function {body_fn}, "
-                f"expected {len(res[-1])}"
-            )
-            res.append(r)
-        else:
-            assert isinstance(r, torch.Tensor), (
-                f"Unexpected type {r} for function {body_fn}, "
-                f"it must be a tuple or a Tensor."
-            )
-            assert not res or len(res[-1]) == 1, (
-                f"Unexpected number of results {len(r)} for function {body_fn}, "
-                f"expected {len(res[-1])}"
-            )
-            res.append((r,))
-
-    if not res:
-        return torch.empty(tuple(), dtype=torch.float32, device=args[0].device)
-    if len(res) == 1:
-        final = res[0]
-    else:
-        n_res = len(res[0])
-        final = [
-            torch.cat(
-                [r[i] for r in res],
-                dim=(
-                    0 if reduction_dim is None or i >= len(reduction_dim) else reduction_dim[i]
-                ),
-            )
-            for i in range(n_res)
-        ]
-    return tuple(final) if len(final) > 1 else final[0]
-
-
-def make_custom_loop_for(
-    n_iter: torch.Tensor,
-    body_fn: Callable,
-    reduction_dim: Optional[Sequence[int]],
-    args: Sequence[torch.Tensor],
-    body_gm: Optional[torch.fx.GraphModule] = None,
-    body_mutated_inputs: Optional[List[Any]] = None,
-    body_outputs: Optional[List[Any]] = None,
-) -> Tuple[str, torch.library.CustomOpDef]:
-    """
-    Defines a custom operator for a loop in order to avoid
-    :func:`torch.export.export` digging into it.
-    It registers the custom op and a custom conversion
-    to ONNX.
-
-    :param n_iter: number of iterations defined by a tensor of no dimension
-    :param body_fn: the loop body defined as a function
-    :param reduction_dim: dimension used to concatenated the results
-    :param args: list of tensors, input to the body
-    :param body_gm: torch.fx.GraphModule equivalent to *body_gm*
-    :param body_mutated_inputs: inputs to *body_gm*
-    :param body_outputs: outputs to *body_gm*
-    :return: a name and the custom op definition, the name
-        is used to cache the custom op
-    """
-    assert body_gm is not None, "body_gm cannot be None"
-    assert body_mutated_inputs is not None, "body_mutated_inputs cannot be None"
-    assert body_outputs is not None, "body_outputs cannot be None"
-
-    srank = "_".join("x".join(map(str, s.shape)) for s in body_outputs)
-    sred = "x".join(map(str, reduction_dim)) if reduction_dim else ""
-    full_name = (
-        body_fn.__qualname__.replace("<locals>", "L")
-        .replace("<lambda>", "l")
-        .replace(".", "_")
-    )
-    name = f"loop_for_onnx_{full_name}_{srank}_{sred}"
-
-    schema = "(str body_fn, Tensor n_iter, Tensor[] body_inputs) -> Tensor"
-    if len(body_outputs) > 1:
-        schema += "[]"
-    custom_def = torch.library.CustomOpDef("onnx_higher_ops", "loop_for", schema, body_fn)
-    custom_def.register_kernel("cpu")(body_fn)
-
-    custom_def._abstract_fn = lambda _fn_id, *_args, _o=body_outputs: (
-        tuple([torch.empty_like(s) for s in _o]) if len(_o) > 1 else torch.empty_like(_o[0])
-    )
-    return name, custom_def
-
-
-def loop_for(
-    n_iter: Union[torch.SymInt, torch.Tensor],
-    body_fn: Callable[..., Tuple[torch.Tensor]],
-    args: Sequence[torch.Tensor],
-    reduction_dim: Optional[Sequence[int]] = None,
-) -> Tuple[torch.Tensor, ...]:
-    """
-    High operators used to easily export a loop in ONNX.
-    Does not fully work with :func:`torch.export.export`,
-    it does replaces a custom op with a loop operator afterwards.
-    Every iteration produces tensors, all of them are gathered
-    into lists, all these lists are concatenated into tensors.
-
-    :param n_iter: number of iterations, it can be fixed on
-        variable, in that case it should a tensor with no dimension
-    :param body_fn: function body, takes only tensors and returns
-        only tensors, the first argument is the iteration number
-        in a tensor with no dimension, all the others
-        are not changed during the loop
-    :param args: the available tensors at every loop
-    :param reduction_dim: the loop aggregated the results into list,
-        one of each output, each of them is concatenated into one
-        tensor along one dimension, by default, it is the first
-        dimension, but it can be defined otherwise
-    """
-    assert args, "The function should have at least one arg."
-    assert (
-        isinstance(n_iter, torch.Tensor)
-        and n_iter.dtype == torch.int64
-        and len(n_iter.shape) == 0
-    ), f"Only a tensor for one int64 is allowed for n_iter but it equal to {n_iter}."
-    if is_exporting():
-        from torch.fx.experimental.proxy_tensor import _CURRENT_MAKE_FX_TRACER
-
-        # tracer = _CURRENT_MAKE_FX_TRACER.fx_tracer
-        root = _CURRENT_MAKE_FX_TRACER.fx_tracer.root
-        # graph = _CURRENT_MAKE_FX_TRACER.fx_tracer.graph
-
-        body_gm: torch.fx.GraphModule = materialize_as_graph(
-            body_fn, (torch.tensor(0, dtype=torch.int64), *args)
-        )
-        (
-            _1,
-            _2,
-            _3,
-            body_mutated_inputs,
-            body_outputs,
-        ) = check_input_alias_and_mutation_return_outputs(body_gm)
-        name, _custom_ops = make_custom_loop_for(
-            n_iter,
-            body_fn,
-            reduction_dim,
-            args,
-            body_gm=body_gm,
-            body_mutated_inputs=body_mutated_inputs,
-            body_outputs=body_outputs,
-        )
-        root.register_module(name, body_gm)
-        # body_graph = _maybe_reenter_make_fx(body_fn)(n_iter, *args)
-        return torch.ops.onnx_higher_ops.loop_for(name, n_iter, args)
-
-    return _loop_for_fn(n_iter, body_fn, reduction_dim, args)
-
-
-"""
-        proxy_mode.tracer.root.register_module(cond_graph_name, cond_graph)
-        proxy_mode.tracer.root.register_module(body_graph_name, body_graph)
-
-        args = (cond_graph, body_graph, carried_inputs, additional_inputs)
-
-        proxy_args = pytree.tree_map(proxy_mode.tracer.unwrap_proxy, args)
-
-        out_proxy = proxy_mode.tracer.create_proxy(
-            "call_function", op, proxy_args, {}, name=op._name
-        )
-
-        out = op(
-            cond_graph, body_graph, unspecialized_carried_inputs, additional_inputs
-        )
-        return track_tensor_tree(
-            out, out_proxy, constant=None, tracer=proxy_mode.tracer
-        )
-"""