assert

Author: xadupre

onnx_diagnostic/helpers/torch_helper.py

@@ -693,9 +693,13 @@ def to_any(value: Any, to_value: Union[torch.dtype, torch.device, str]) -> Any:
     """Applies torch.to if applicable. Goes recursively."""
     if isinstance(value, (torch.nn.Module, torch.Tensor)):
         if (
-            isinstance(to_value, torch.dtype)
-            or to_value in {"float16", "bfloat16", "float32", "float64"}
-        ) and value.dtype in {torch.int32, torch.int64, torch.int8, torch.int16}:
+            (
+                isinstance(to_value, torch.dtype)
+                or to_value in {"float16", "bfloat16", "float32", "float64"}
+            )
+            and hasattr(value, "dtype")
+            and value.dtype in {torch.int32, torch.int64, torch.int8, torch.int16}
+        ):
             # int vector should not be changed.
             return value
         return value.to(to_value)

onnx_diagnostic/torch_export_patches/onnx_export_errors.py

@@ -188,9 +188,9 @@ def torch_export_patches(
     if rewrite:
         from .patch_module import torch_export_rewrite
 
-        with torch_export_rewrite(  # type: ignore[var-annotated]
+        with torch_export_rewrite(
             rewrite=rewrite, dump_rewriting=dump_rewriting, verbose=verbose
-        ), torch_export_patches(
+        ), torch_export_patches(  # type: ignore[var-annotated]
             patch_sympy=patch_sympy,
             patch_torch=patch_torch,
             patch_transformers=patch_transformers,

onnx_diagnostic/torch_export_patches/patch_module.py

@@ -177,16 +177,21 @@ def _rewrite_if(
             else_ret = else_exprs[0]
             then_exprs = [n for n in node.body if not isinstance(n, ast.Return)]
             else_exprs = [n for n in node.orelse if not isinstance(n, ast.Return)]
-            assert type(then_ret.value) is type(else_ret.value), (
-                f"Inconsistencies return then value={then_ret.value}, "
-                f"else value={else_ret.value}"
+            is_tuple_or_list = (
+                isinstance(then_ret, (ast.Tuple, ast.List)),
+                isinstance(else_ret, (ast.Tuple, ast.List)),
             )
-            if isinstance(then_ret.value, (ast.Tuple, ast.list)):
-                assert len(then_ret.value.elts) == len(else_ret.value.elts), (
+            assert len(set(is_tuple_or_list)) == 1, (
+                f"is_tuple_or_list={is_tuple_or_list}, inconsistencies return "
+                f"then value={then_ret}, "
+                f"else value={else_ret}"
+            )
+            if is_tuple_or_list[0]:
+                assert len(then_ret.elts) == len(else_ret.elts), (
                     f"Unexpected number of elements on both branches, "
-                    f"then:{then_ret.value.elts}, else:{else_ret.value.elts}"
+                    f"then:{then_ret.elts}, else:{else_ret.elts}"
                 )
-                n_returned_values = len(then_ret.value.elts)
+                n_returned_values = len(then_ret.elts)
             else:
                 n_returned_values = 0
         else: