fix test

Author: xadupre

_doc/status/index.rst

@@ -11,6 +11,7 @@ what works and what does not with :func:`torch.export.export`.
     exported_program_dynamic
     exporter_dynamic
     patches_coverage
+    patches_diff
 
 Examples checking about dynamic dimensions:
 

_doc/status/patches_diff.rst

@@ -0,0 +1,72 @@
+.. _l-patch-diff:
+
+============
+Patches Diff
+============
+
+Patches are not always needed to export a LLM.
+Most of the time, only serialization function are needed to export
+a LLM with cache (``DynamicCache``, ...).
+Function :func:`register_additional_serialization_functions
+<onnx_diagnostic.torch_export_patches.register_additional_serialization_functions>`
+is enough in many cases.
+
+.. code-block:: python
+
+    import torch
+    from onnx_diagnostic.torch_export_patches import register_additional_serialization_functions
+
+    with register_additional_serialization_functions(patch_transformers=True):
+        ep = torch.export.export(...)
+
+Function :func:`torch_export_patches
+<onnx_diagnostic.torch_export_patches.torch_export_patches>`
+helps fixing some issues for many models.
+
+.. code-block:: python
+
+    import torch
+    from onnx_diagnostic.torch_export_patches import torch_export_patches
+
+    with torch_export_patches(patch_transformers=True):
+        ep = torch.export.export(...)
+
+Class :class:`PatchDetails <onnx_diagnostic.torch_export_patches.patch_details.PatchDetails>`
+gives an example on how to retrieve the list of involded patches for a specific model.
+Those patches belongs to the following list which depends on transformers and
+pytorch versions.
+
+.. runpython::
+    :showcode:
+
+    import torch
+    import transformers
+
+    print(torch.__version__, transformers.__version__)
+
+Those two versions leads to the following list of patches.
+
+.. runpython::
+    :showcode:
+    :rst:
+
+    from onnx_diagnostic.torch_export_patches.patch_details import PatchDetails
+    from onnx_diagnostic.torch_export_patches import torch_export_patches
+
+    details = PatchDetails()
+    with torch_export_patches(
+        patch_transformers=True,
+        patch_torch=True,
+        patch_diffusers=True,
+        patch_details=details,
+    ):
+        pass
+    for patch in details.patched:
+        if patch.function_to_patch == patch.patch:
+            continue
+        rst = patch.format_diff(format="rst")
+        print()
+        print()
+        print(rst)
+        print()
+        print()

_unittests/ut_torch_export_patches/test_patch_details.py

@@ -5,7 +5,6 @@
 from onnx_diagnostic.torch_export_patches import torch_export_patches
 from onnx_diagnostic.torch_export_patches.patch_inputs import use_dyn_not_str
 from onnx_diagnostic.torch_export_patches.patch_details import PatchDetails, PatchInfo
-from onnx_diagnostic.torch_export_patches.patches.patch_transformers import patched_eager_mask
 from onnx_diagnostic.torch_models.hghub import get_untrained_model_with_inputs
 
 
@@ -35,6 +34,10 @@ def test_patch_details(self):
 
     @requires_transformers("4.55")
     def test_patch_diff(self):
+        from onnx_diagnostic.torch_export_patches.patches.patch_transformers import (
+            patched_eager_mask,
+        )
+
         eager_mask = transformers.masking_utils.eager_mask
         self.assertEqual(eager_mask.__name__, "eager_mask")
         self.assertEqual(patched_eager_mask.__name__, "patched_eager_mask")

onnx_diagnostic/torch_export_patches/patch_details.py

@@ -1,5 +1,6 @@
 import difflib
 import inspect
+import pprint
 import re
 import textwrap
 from typing import Any, Dict, Callable, List, Optional, Tuple, Union
@@ -111,9 +112,13 @@ def format_diff(self, format: str = "raw") -> str:
 
             import transformers
             import onnx_diagnostic.torch_export_patches.patches.patch_transformers as ptr
-            from onnx_diagnostic.torch_export_patches.patch_details import Patchinfo
+            from onnx_diagnostic.torch_export_patches.patch_details import PatchInfo
+            from onnx_diagnostic.torch_export_patches.patches.patch_transformers import (
+                patched_eager_mask,
+            )
 
-            diff = Patchinfo(eager_mask, patched_eager_mask).format_diff(format="rst")
+            eager_mask = transformers.masking_utils.eager_mask
+            diff = PatchInfo(eager_mask, patched_eager_mask).format_diff(format="rst")
             print(diff)
         """
         diff = self.make_diff()
@@ -146,7 +151,8 @@ class PatchDetails:
     """
     This class is used to store patching information.
     This helps understanding which rewriting was applied to which
-    method of functions.
+    method of functions. Page :ref:`l-patch-diff` contains all the
+    diff for all the implemented patches.
 
     .. runpython::
         :showcode:
@@ -227,6 +233,7 @@ def patches_involded_in_graph(
             node_stack.append((node, stack))
 
         patch_node = []
+        patched_nodes = set()
         for patch, _f, source, interval in patches:
             exp = 'File "([^"]*?%s[^"]+?)", line (\\d+)' % cst
             reg = re.compile(exp)
@@ -246,6 +253,15 @@ def patches_involded_in_graph(
                         and self.matching_pair(patch, node)
                     ):
                         patch_node.append((patch, node))
+                        patched_nodes.add(id(node))
+
+        # checks all patches were discovered
+        for node, _ in node_stack:
+            assert id(node) in patched_nodes, (
+                f"One node was patched but no patch was found:\n"
+                f"node: {node.target}({','.join(map(str, node.args))}) -> {node.name}"
+                f"\n--\n{pprint.pformat(node.meta)}"
+            )
 
         res = {}
         for patch, node in patch_node: