Make the code more agnostic to custom classes

Author: xadupre

onnx_diagnostic/export/dynamic_shapes.py

@@ -313,6 +313,24 @@ def guess_dynamic_shape_object(self, *objs: Any, msg: Optional[Callable] = None)
                 shapes[i] = self.guess_dynamic_shape_object(*[o[i] for o in objs])
             return shapes
 
+        if obj.__class__ in torch.utils._pytree.SUPPORTED_NODES:
+            kcl = set(o.__class__ for o in objs)
+            assert len(kcl) == 1, (
+                f"All instances of argument {i} are not of the same class but {kcl}, "
+                f"types should be the same."
+            )
+            col_args = [torch.utils._pytree.tree_flatten(o) for o in objs]
+            kc = set(len(col_args) for o in objs)
+            assert len(kc) == 1, (
+                f"All instances of type {kcl.pop()} are not serialized into the same number "
+                f"of arguments, it should be the same."
+            )
+            values = []
+            for i in range(kc.pop()):
+                values.append(self.guess_dynamic_dimensions(*[ca[i] for ca in col_args]))
+            return values
+
+        # In case DynamicCache is not registered.
         if obj.__class__.__name__ == "DynamicCache":
             kc = set(len(o.key_cache) for o in objs)
             assert (

onnx_diagnostic/ext_test_case.py

@@ -898,7 +898,7 @@ def assertEqualAny(
                 for e, g in zip(expected, value):
                     self.assertEqualAny(e, g, msg=msg, atol=atol, rtol=rtol)
         elif expected.__class__.__name__ == "DynamicCache":
-            atts = {"key_cache", "value_cache"}
+            atts = ["key_cache", "value_cache"]
             self.assertEqualAny(
                 {k: expected.__dict__.get(k, None) for k in atts},
                 {k: value.__dict__.get(k, None) for k in atts},

onnx_diagnostic/helpers/helper.py

@@ -443,6 +443,17 @@ def string_type(
     if ignore:
         return f"{obj.__class__.__name__}(...)"
 
+    if obj.__class__ in torch.utils._pytree.SUPPORTED_NODES:
+        args, _spec = torch.utils._pytree.tree_flatten(obj)
+        att = string_type(
+            args,
+            with_shape=with_shape,
+            with_min_max=with_min_max,
+            with_device=with_device,
+            limit=limit,
+        )
+        return f"{obj.__class__.__name__}({att})"
+
     raise AssertionError(f"Unsupported type {type(obj).__name__!r} - {type(obj)}")
 
 
@@ -1125,6 +1136,29 @@ def max_diff(
             flatten=flatten,
         )
 
+    if expected.__class__ in torch.utils._pytree.SUPPORTED_NODES:
+        if got.__class__ not in torch.utils._pytree.SUPPORTED_NODES:
+            return dict(abs=np.inf, rel=np.inf, sum=np.inf, n=np.inf, dnan=np.inf)
+        if verbose >= 6:
+            print(
+                f"[max_diff] {expected.__class__.__name__}: "
+                f"{string_type(expected)} ? {string_type(got)}"
+            )
+        expected_args, _spec = torch.utils._pytree.tree_flatten(expected)
+        got_args, _spec = torch.utils._pytree.tree_flatten(got)
+        return max_diff(
+            expected_args,
+            got_args,
+            level=level,
+            flatten=flatten,
+            debug_info=debug_info,
+            begin=begin,
+            end=end,
+            _index=_index,
+            verbose=verbose,
+        )
+
+    # backup function in case pytorch does not know how to serialize.
     if expected.__class__.__name__ == "DynamicCache":
         if got.__class__.__name__ == "DynamicCache":
             if verbose >= 6:

onnx_diagnostic/helpers/torch_test_helper.py

@@ -305,7 +305,7 @@ def forward(self, input_ids):
 
 def to_any(value: Any, to_value: Union[torch.dtype, torch.device]) -> Any:
     """
-    Applies torch.to is applicables.
+    Applies torch.to is applicable.
     Goes recursively.
     """
     if isinstance(value, (torch.nn.Module, torch.Tensor)):
@@ -329,6 +329,10 @@ def to_any(value: Any, to_value: Union[torch.dtype, torch.device]) -> Any:
                 )
             )
         )
+    if value.__class__ in torch.utils._pytree.SUPPORTED_NODES:
+        args, spec = torch.utils._pytree.tree_flatten(value)
+        new_args = to_any(args, to_value)
+        return torch.utils._pytree.tree_unflatten(new_args, spec)
 
     assert not isinstance(value, Iterable), f"Unsupported type {type(value)}"
     return value
@@ -361,6 +365,11 @@ def torch_deepcopy(value: Any) -> Any:
             torch_deepcopy(value.self_attention_cache),
             torch_deepcopy(value.cross_attention_cache),
         )
+    if value.__class__ in torch.utils._pytree.SUPPORTED_NODES:
+        args, spec = torch.utils._pytree.tree_flatten(value)
+        new_args = torch_deepcopy(args)
+        return torch.utils._pytree.tree_unflatten(new_args, spec)
+
     # We should have a code using serialization, deserialization assuming a model
     # cannot be exported without them.
     raise NotImplementedError(f"torch_deepcopy not implemented for type {type(value)}")

onnx_diagnostic/torch_export_patches/patch_inputs.py

@@ -39,6 +39,11 @@ def _make_shape(subset: Dict, cls: type, value: Any) -> Any:
             break
         new_shape = [[axes for i in range(cache_length)], [axes for i in range(cache_length)]]
         return new_shape
+    if value.__class__ in torch.utils._pytree.SUPPORTED_NODES:
+        raise NotImplementedError(
+            f"_make_shape not implemented for registered class={cls}, "
+            f"subset={subset}, value={string_type(value)}"
+        )
     raise NotImplementedError(
         f"_make_shape not implemented for cls={cls}, "
         f"subset={subset}, value={string_type(value)}"