fix issues

Author: xadupre

_unittests/ut_export/test_dynamic_shapes.py

@@ -5,6 +5,7 @@
 from onnx_diagnostic.helpers.cache_helper import make_dynamic_cache
 from onnx_diagnostic.export import ModelInputs
 from onnx_diagnostic.export.dynamic_shapes import CoupleInputsDynamicShapes
+from onnx_diagnostic.torch_export_patches import bypass_export_some_errors
 
 
 class TestDynamicShapes(ExtTestCase):
@@ -529,25 +530,26 @@ def test_couple_input_ds_cache(self):
 
         kwargs = {"A": T3x4, "B": (T3x1, cache)}
         Cls = CoupleInputsDynamicShapes
-        self.assertEqual(
-            [],
-            Cls(
-                (),
-                kwargs,
-                {"A": ds_batch, "B": (ds_batch, [ds_batch, ds_batch, ds_batch, ds_batch])},
-            ).invalid_paths(),
-        )
-        self.assertEqual(
-            [("B", 1, "DynamicCache", 1, "[2]"), ("B", 1, "DynamicCache", 3, "[2]")],
-            Cls(
-                (),
-                kwargs,
-                {
-                    "A": ds_batch,
-                    "B": (ds_batch, [ds_batch, ds_batch_seq, ds_batch, ds_batch_seq]),
-                },
-            ).invalid_paths(),
-        )
+        with bypass_export_some_errors(patch_transformers=True):
+            self.assertEqual(
+                [],
+                Cls(
+                    (),
+                    kwargs,
+                    {"A": ds_batch, "B": (ds_batch, [ds_batch, ds_batch, ds_batch, ds_batch])},
+                ).invalid_paths(),
+            )
+            self.assertEqual(
+                [("B", 1, "DynamicCache", 1, "[2]"), ("B", 1, "DynamicCache", 3, "[2]")],
+                Cls(
+                    (),
+                    kwargs,
+                    {
+                        "A": ds_batch,
+                        "B": (ds_batch, [ds_batch, ds_batch_seq, ds_batch, ds_batch_seq]),
+                    },
+                ).invalid_paths(),
+            )
 
 
 if __name__ == "__main__":

onnx_diagnostic/export/dynamic_shapes.py

@@ -541,6 +541,11 @@ def _valid_shapes(
                             yield path
             else:
                 # A custom class.
+                assert inputs.__class__ in torch.utils._pytree.SUPPORTED_NODES, (
+                    f"Class {inputs.__class__.__name__!r} was not registered using "
+                    f"torch.utils._pytree.register_pytree_node, it is not possible to "
+                    f"map this class with the given dynamic shapes."
+                )
                 flat, _spec = torch.utils._pytree.tree_flatten(inputs)
                 for path in cls._valid_shapes(
                     flat, ds, prefix=(*prefix, inputs.__class__.__name__)