fix issues

Author: xadupre

_doc/conf.py

@@ -90,7 +90,8 @@ def linkcode_resolve(domain, info):
         "https://sdpython.github.io/doc/experimental-experiment/dev/",
         None,
     ),
-    "diffusers": ("https://huggingface.co/docs/diffusers/index", None),
+    # Not a sphinx documentation
+    # "diffusers": ("https://huggingface.co/docs/diffusers/index", None),
     "matplotlib": ("https://matplotlib.org/stable/", None),
     "numpy": ("https://numpy.org/doc/stable", None),
     "onnx": ("https://onnx.ai/onnx/", None),
@@ -105,7 +106,8 @@ def linkcode_resolve(domain, info):
     "sklearn": ("https://scikit-learn.org/stable/", None),
     "skl2onnx": ("https://onnx.ai/sklearn-onnx/", None),
     "torch": ("https://pytorch.org/docs/main/", None),
-    "transformers": ("https://huggingface.co/docs/transformers/index", None),
+    # Not a sphinx documentation
+    # "transformers": ("https://huggingface.co/docs/transformers/index", None),
 }
 
 # Check intersphinx reference targets exist
@@ -118,6 +120,7 @@ def linkcode_resolve(domain, info):
     ("py:class", "True"),
     ("py:class", "Argument"),
     ("py:class", "default=sklearn.utils.metadata_routing.UNCHANGED"),
+    ("py:class", "diffusers.models.unets.unet_2d_condition.UNet2DConditionOutput"),
     ("py:class", "ModelProto"),
     ("py:class", "Model"),
     ("py:class", "Module"),

_doc/patches.rst

@@ -113,7 +113,7 @@ Here is the list of supported caches:
 
     import onnx_diagnostic.torch_export_patches.onnx_export_serialization as p
 
-    print("\n".join(sorted(p.serialization_functions())))
+    print("\n".join(sorted(t.__name__ for t in p.serialization_functions())))
 
 .. _l-control-flow-rewriting:
 

_doc/status/patches_coverage.rst

@@ -14,7 +14,7 @@ The following code shows the list of serialized classes in transformers.
 
     import onnx_diagnostic.torch_export_patches.onnx_export_serialization as p
 
-    print('\n'.join(sorted(p.serialization_functions())))
+    print('\n'.join(sorted(t.__name__ for t in p.serialization_functions())))
 
 Patched Classes
 ===============

_unittests/ut_torch_export_patches/test_patch_serialization.py

@@ -1,7 +1,12 @@
 import unittest
 import torch
 from transformers.modeling_outputs import BaseModelOutput
-from onnx_diagnostic.ext_test_case import ExtTestCase, ignore_warnings, requires_torch
+from onnx_diagnostic.ext_test_case import (
+    ExtTestCase,
+    ignore_warnings,
+    requires_torch,
+    requires_diffusers,
+)
 from onnx_diagnostic.helpers.cache_helper import (
     make_encoder_decoder_cache,
     make_dynamic_cache,
@@ -212,6 +217,56 @@ def test_sliding_window_cache_flatten(self):
                 self.string_type(cache2, with_shape=True, with_min_max=True),
             )
 
+    @ignore_warnings(UserWarning)
+    @requires_diffusers("0.30")
+    def test_unet_2d_condition_output(self):
+        from diffusers.models.unets.unet_2d_condition import UNet2DConditionOutput
+
+        bo = UNet2DConditionOutput(sample=torch.rand((4, 4, 4)))
+        self.assertEqual(bo.__class__.__name__, "UNet2DConditionOutput")
+        with torch_export_patches():
+            # internal function
+            bo2 = torch_deepcopy([bo])
+            self.assertIsInstance(bo2, list)
+            self.assertEqual(bo2[0].__class__.__name__, "UNet2DConditionOutput")
+            self.assertEqualAny([bo], bo2)
+            self.assertEqual(
+                "UNet2DConditionOutput(sample:T1s4x4x4)",
+                self.string_type(bo, with_shape=True),
+            )
+
+            # serialization
+            flat, _spec = torch.utils._pytree.tree_flatten(bo)
+            self.assertEqual(
+                "#1[T1s4x4x4]",
+                self.string_type(flat, with_shape=True),
+            )
+            bo2 = torch.utils._pytree.tree_unflatten(flat, _spec)
+            self.assertEqual(
+                self.string_type(bo, with_shape=True, with_min_max=True),
+                self.string_type(bo2, with_shape=True, with_min_max=True),
+            )
+
+            # flatten_unflatten
+            flat, _spec = torch.utils._pytree.tree_flatten(bo)
+            unflat = flatten_unflatten_for_dynamic_shapes(bo, use_dict=True)
+            self.assertIsInstance(unflat, dict)
+            self.assertEqual(list(unflat), ["sample"])
+
+        # export
+        class Model(torch.nn.Module):
+            def forward(self, cache):
+                return cache.sample[0]
+
+        bo = UNet2DConditionOutput(sample=torch.rand((4, 4, 4)))
+        model = Model()
+        model(bo)
+        DYN = torch.export.Dim.DYNAMIC
+        ds = [{0: DYN}]
+
+        with torch_export_patches():
+            torch.export.export(model, (bo,), dynamic_shapes=(ds,))
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

onnx_diagnostic/helpers/config_helper.py

@@ -71,7 +71,7 @@ def _pick(config, *atts, exceptions: Optional[Dict[str, Callable]] = None):
 
 def pick(config, name: str, default_value: Any) -> Any:
     """
-    Returns the vlaue of a attribute if config has it
+    Returns the value of a attribute if config has it
     otherwise the default value.
     """
     if not config:

onnx_diagnostic/torch_export_patches/onnx_export_serialization.py

@@ -35,7 +35,7 @@
 
 
 PATCH_OF_PATCHES: Set[Any] = set()
-WRONG_REGISTRATIONS: Dict[str, str] = {
+WRONG_REGISTRATIONS: Dict[str, Optional[str]] = {
     DynamicCache: "4.50",
     BaseModelOutput: None,
     UNet2DConditionOutput: None,
@@ -125,7 +125,7 @@ def register_cache_serialization(verbose: int = 0) -> Dict[str, bool]:
         ):
             assert cls in registration_functions, (
                 f"{cls} has no registration functions mapped to it, "
-                f"available {sorted(registration_functions)}"
+                f"available options are {list(registration_functions)}"
             )
             if verbose:
                 print(
@@ -146,7 +146,7 @@ def register_cache_serialization(verbose: int = 0) -> Dict[str, bool]:
     return done
 
 
-def serialization_functions(verbose: int = 0) -> Dict[type, Union[Callable[[], bool], int]]:
+def serialization_functions(verbose: int = 0) -> Dict[type, Union[Callable[[int], bool], int]]:
     """Returns the list of serialization functions."""
     transformers_classes = {
         DynamicCache: lambda verbose=verbose: register_class_serialization(