mypy

Author: xadupre

onnx_diagnostic/export/dynamic_shapes.py

@@ -297,14 +297,14 @@ def _build_new_tensor(self, tensor: torch.Tensor, new_shape: Tuple[int, ...]):
                     tuple(alt_shape), dtype=tensor.dtype, device=tensor.device
                 )
                 mind = min(d0, d1)
-                indices = [slice(None) for _ in range(rank)]
+                indices: List[Union[slice, int]] = [slice(None) for _ in range(rank)]
                 indices[i] = slice(0, mind)
                 ind = tuple(indices)
                 new_tensor[ind] = tensor[ind]
                 if d1 > mind:
                     for k in range(d1 - mind):
-                        indices0 = [slice(None) for _ in range(rank)]
-                        indices1 = [slice(None) for _ in range(rank)]
+                        indices0: List[Union[slice, int]] = [slice(None) for _ in range(rank)]
+                        indices1: List[Union[slice, int]] = [slice(None) for _ in range(rank)]
                         indices1[i] = mind + k
                         indices0[i] = k % mind
                         new_tensor[tuple(indices1)] = tensor[tuple(indices0)]

onnx_diagnostic/export/validate.py

@@ -9,7 +9,7 @@
 
 
 def compare_modules(
-    modep: torch.export.ExportedProgram,
+    modep: torch.nn.Module,
     mod: Optional[torch.nn.Module] = None,
     args: Optional[Tuple[Any, ...]] = None,
     kwargs: Optional[Dict[str, Any]] = None,
@@ -18,7 +18,7 @@ def compare_modules(
     verbose: int = 0,
     atol: float = 1e-2,
     rtol: float = 1e-1,
-) -> List[Dict[str, Any]]:
+) -> Dict[str, Any]:
     """
     Compares two torch modules, usually one coming from an exported program,
     the other being the origin model.
@@ -150,8 +150,8 @@ def validate_ep(
     values = [_[1] for _ in items]
     all_vals = list(itertools.product(*values))
     cpl = CoupleInputsDynamicShapes(
-        args,
-        kwargs,
+        args or (),
+        kwargs or {},
         dynamic_shapes,
         args_names=(
             list(inspect.signature(modep.forward).parameters) if args and kwargs else None

onnx_diagnostic/torch_models/hghub/hub_api.py

@@ -60,7 +60,7 @@ def get_pretrained_config(
     )
 
 
-def get_model_info(model_id) -> str:
+def get_model_info(model_id) -> Any:
     """Returns the model info for a model_id."""
     return model_info(model_id)
 
@@ -220,7 +220,11 @@ def enumerate_model_list(
                             m.trending_score or "",
                             m.private or "",
                             m.gated or "",
-                            ("|".join(m.tags)).replace(",", "_").replace(" ", "_"),
+                            (
+                                ("|".join(m.tags)).replace(",", "_").replace(" ", "_")
+                                if m.tags
+                                else ""
+                            ),
                         ],
                     )
                 )