onnx-diagnostic 0.8.10__py3-none-any.whl → 0.8.11__py3-none-any.whl

onnx_diagnostic/ext_test_case.py

@@ -780,7 +780,7 @@ class ExtTestCase(unittest.TestCase):
 
     @property
     def verbose(self) -> int:
-        "Returns the the value of environment variable ``VERBOSE``."
+        "Returns the value of environment variable ``VERBOSE``."
         return int(os.environ.get("VERBOSE", "0"))
 
     @classmethod
@@ -1028,6 +1028,19 @@ class ExtTestCase(unittest.TestCase):
                 rtol=rtol,
                 msg=msg,
             )
+        elif expected.__class__.__name__ == "BaseModelOutputWithPooling":
+            if expected.__class__.__name__ == value.__class__.__name__:
+                self.assertEqual(len(expected), len(value), msg=msg)
+                self.assertEqual(list(expected), list(value), msg=msg)  # checks the order
+                self.assertEqualAny(
+                    {k: v for k, v in expected.items()},  # noqa: C416
+                    {k: v for k, v in value.items()},  # noqa: C416
+                    atol=atol,
+                    rtol=rtol,
+                    msg=msg,
+                )
+            else:
+                self.assertEqualArray(expected.last_hidden_state, value)
         elif isinstance(expected, (tuple, list, dict)):
             self.assertIsInstance(value, type(expected), msg=msg)
             self.assertEqual(len(expected), len(value), msg=msg)
@@ -1043,24 +1056,28 @@ class ExtTestCase(unittest.TestCase):
             "SlidingWindowCache",
             "HybridCache",
         ):
+            from .helpers.cache_helper import CacheKeyValue
+
             self.assertEqual(type(expected), type(value), msg=msg)
-            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},
-                atol=atol,
-                rtol=rtol,
-            )
+            self.assertEqualAny(CacheKeyValue(expected), CacheKeyValue(value))
         elif expected.__class__.__name__ == "StaticCache":
+            from .helpers.cache_helper import CacheKeyValue
+
             self.assertEqual(type(expected), type(value), msg=msg)
             self.assertEqual(expected.max_cache_len, value.max_cache_len)
-            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},
-                atol=atol,
-                rtol=rtol,
-            )
+            self.assertEqualAny(CacheKeyValue(expected), CacheKeyValue(value))
+        elif expected.__class__.__name__ == "CacheKeyValue":
+            self.assertEqual(type(expected), type(value), msg=msg)
+            if expected.cls_layers is None:
+                self.assertEqual(expected.cls_layers, value.cls_layers)
+            else:
+                self.assertEqualAny(
+                    [cls.__name__ for cls in expected.cls_layers],
+                    [cls.__name__ for cls in value.cls_layers],
+                    msg=msg,
+                )
+            self.assertEqualAny(expected.key_cache, value.key_cache, msg=msg)
+            self.assertEqualAny(expected.value_cache, value.value_cache, msg=msg)
         elif expected.__class__.__name__ == "EncoderDecoderCache":
             self.assertEqual(type(expected), type(value), msg=msg)
             atts = ["self_attention_cache", "cross_attention_cache"]

onnx_diagnostic/helpers/args_helper.py

@@ -105,6 +105,7 @@ def get_parsed_args(
                 default=tries,
             )
     for k, v in kwargs.items():
+        assert isinstance(v, tuple)  # type
         parser.add_argument(
             f"--{k}",
             help=f"{v[1]}, default is {v[0]}",

onnx_diagnostic/helpers/bench_run.py

@@ -11,7 +11,6 @@ from argparse import Namespace
 from datetime import datetime
 from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
 
-
 _DEFAULT_STRING_LIMIT = 2000
 
 

onnx_diagnostic/helpers/cache_helper.py

@@ -90,7 +90,7 @@ def flatten_unflatten_for_dynamic_shapes(
         the context gives the dictionary keys but it is not expressed
         in the dynamic shapes, these specifications seems to be different
         for the strict and non strict mode. It also preserves tuple.
-    :param change_function: to modifies the tensor in the structure itself,
+    :param change_function: to modify the tensor in the structure itself,
         like replace them by a shape
     :return: the serialized object
     """
@@ -110,7 +110,7 @@ def flatten_unflatten_for_dynamic_shapes(
         start = end
     if use_dict:
         if spec.type is dict:
-            # This a dictionary.
+            # This is a dictionary.
             return dict(zip(spec.context, subtrees))
         if spec.type is tuple:
             return tuple(subtrees)
@@ -348,6 +348,7 @@ else:
 def make_static_cache(
     key_value_pairs: Union[List[torch.Tensor], List[Tuple[torch.Tensor, torch.Tensor]]],
     max_cache_len: Optional[int] = None,
+    cls_layers: Optional[Union[str, List[type]]] = None,
 ) -> transformers.cache_utils.DynamicCache:
     """
     Creates an instance of :class:`transformers.cache_utils.StaticCache`.
@@ -379,6 +380,9 @@ def make_static_cache(
         )
         print(string_type(past_key_values, with_shape=True))
     """
+    assert not cls_layers or set(cls_layers) == {
+        transformers.cache_utils.StaticLayer
+    }, f"Not implemented when cls_layers={cls_layers!r}"
     key_value_pairs = _preprocess_key_value_pairs(key_value_pairs)
 
     class _config:
@@ -583,13 +587,9 @@ if hasattr(transformers.cache_utils, "SlidingWindowCache"):
         )
         return finalize_cache(cache)
 
-    def get_make_hybrid_cache():
-        return make_sliding_window_cache
-
 else:
     make_sliding_window_cache = None  # type: ignore[assignment]
 
-
 if hasattr(transformers.cache_utils, "HybridCache"):
 
     def make_hybrid_cache(

onnx_diagnostic/helpers/doc_helper.py

@@ -1,5 +1,5 @@
 import os
-from typing import Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Optional, Tuple
 import onnx
 import onnx.helper as oh
 import torch
@@ -46,10 +46,10 @@ class LayerNormalizationOrt(OpRunKernel):
             f"This kernel implementation only work when only one output "
             f"is required but {node.output} were."
         )
-        self._cache: Dict[Tuple[int, int], onnx.ModelProto] = {}
+        self._cache: Dict[Tuple[int, int], Any] = {}
         self.is_cpu = torch.device("cpu") == self.device
 
-    def _make_model(self, itype: int, rank: int, has_bias: bool) -> onnx.ModelProto:
+    def _make_model(self, itype: int, rank: int, has_bias: bool) -> Any:
         shape = [*["d{i}" for i in range(rank - 1)], "last"]
         layer_model = oh.make_model(
             oh.make_graph(
@@ -88,6 +88,7 @@ class LayerNormalizationOrt(OpRunKernel):
             providers=[provider],
         )
 
+    # pyrefly: ignore[bad-override]
     def run(self, x, scale, bias=None):
         itype = torch_dtype_to_onnx_dtype(x.dtype)
         rank = len(x.shape)
@@ -124,7 +125,7 @@ class MatMulOrt(OpRunKernel):
         self._cache: Dict[Tuple[int, int, int], onnx.ModelProto] = {}
         self.is_cpu = torch.device("cpu") == self.device
 
-    def _make_model(self, itype: int, ranka: int, rankb: int) -> onnx.ModelProto:
+    def _make_model(self, itype: int, ranka: int, rankb: int) -> Any:
         shapea = ["a{i}" for i in range(ranka)]
         shapeb = ["b{i}" for i in range(rankb)]
         shapec = ["c{i}" for i in range(max(ranka, rankb))]
@@ -149,6 +150,7 @@ class MatMulOrt(OpRunKernel):
             providers=[provider],
         )
 
+    # pyrefly: ignore[bad-override]
     def run(self, a, b):
         itype = torch_dtype_to_onnx_dtype(a.dtype)
         ranka, rankb = len(a.shape), len(b.shape)
@@ -159,5 +161,6 @@ class MatMulOrt(OpRunKernel):
         if self.verbose:
             print(f"[MatMulOrt] running on {self._provider!r}")
         feeds = dict(A=a.tensor, B=b.tensor)
+        # pyrefly: ignore[missing-attribute]
         got = sess.run(None, feeds)[0]
         return OpRunTensor(got)

onnx_diagnostic/helpers/graph_helper.py

@@ -36,7 +36,7 @@ class GraphRendering:
         :return: computation order
         """
         assert not ({"If", "Scan", "Loop", "SequenceMap"} & set(n.op_type for n in nodes)), (
-            f"This algorithme is not yet implemented if the sequence contains "
+            f"This algorithm is not yet implemented if the sequence contains "
             f"a control flow, types={sorted(set(n.op_type for n in nodes))}"
         )
         number = {e: start - 1 for e in (existing or [])}  # noqa: C420
@@ -131,14 +131,14 @@ class GraphRendering:
     @property
     def nodes(self) -> List[onnx.NodeProto]:
         "Returns the list of nodes"
-        return (
+        return list(
             self.proto.graph.node
             if isinstance(self.proto, onnx.ModelProto)
             else self.proto.node
         )
 
     @property
-    def start_names(self) -> List[onnx.NodeProto]:
+    def start_names(self) -> List[str]:
         "Returns the list of known names, inputs and initializer"
         graph = self.proto.graph if isinstance(self.proto, onnx.ModelProto) else self.proto
         input_names = (
@@ -151,7 +151,7 @@ class GraphRendering:
             if isinstance(graph, onnx.FunctionProto)
             else [
                 *[i.name for i in graph.initializer],
-                *[i.name for i in graph.sparse_initializer],
+                *[i.values.name for i in graph.sparse_initializer],
             ]
         )
         return [*input_names, *init_names]
@@ -159,7 +159,7 @@ class GraphRendering:
     @property
     def input_names(self) -> List[str]:
         "Returns the list of input names."
-        return (
+        return list(
             self.proto.input
             if isinstance(self.proto, onnx.FunctionProto)
             else [
@@ -173,7 +173,7 @@ class GraphRendering:
     @property
     def output_names(self) -> List[str]:
         "Returns the list of output names."
-        return (
+        return list(
             self.proto.output
             if isinstance(self.proto, onnx.FunctionProto)
             else [

onnx_diagnostic/helpers/log_helper.py

@@ -29,10 +29,10 @@ class CubeViewDef:
     :param order: to reorder key in columns index
     :param key_agg: aggregate according to these columns before
         creating the view
-    :param agg_args: see :meth:`pandas.core.groupby.DataFrameGroupBy.agg`,
+    :param agg_args: see :meth:`pandas.api.typing.DataFrameGroupBy.agg`,
         it can be also a callable to return a different aggregation
         method depending on the column name
-    :param agg_kwargs: see :meth:`pandas.core.groupby.DataFrameGroupBy.agg`
+    :param agg_kwargs: see :meth:`pandas.api.typing.DataFrameGroupBy.agg`
     :param agg_multi: aggregation over multiple columns
     :param ignore_columns: ignore the following columns if known to overload the view
     :param keep_columns_in_index: keeps the columns even if there is only one unique value
@@ -98,7 +98,7 @@ class CubeViewDef:
         agg_args: Union[Sequence[Any], Callable[[str], Any]] = ("sum",),
         agg_kwargs: Optional[Dict[str, Any]] = None,
         agg_multi: Optional[
-            Dict[str, Callable[[pandas.core.groupby.DataFrameGroupBy], pandas.Series]]
+            Dict[str, Callable[[pandas.api.typing.DataFrameGroupBy], pandas.Series]]
         ] = None,
         ignore_columns: Optional[Sequence[str]] = None,
         keep_columns_in_index: Optional[Sequence[str]] = None,
@@ -365,6 +365,7 @@ class CubePlot:
                     # This is very slow
                     # ddd.plot(ax=axs[row, ii],linewidth=3)
                     for jj in range(ddd.shape[1]):
+                        # pyrefly: ignore[bad-index]
                         axs[row, ii].plot(x, ddd.iloc[:, jj], lw=3, label=ddd.columns[jj])
                     axs[row, ii].set_title(f"{c}{title_suffix}")
                     rotate_align(axs[row, ii])
@@ -480,7 +481,9 @@ class CubeLogs:
         elif isinstance(self._data, list) and all(isinstance(r, dict) for r in self._data):
             if verbose:
                 print(f"[CubeLogs.load] load from list of dicts, n={len(self._data)}")
-            self.data = pandas.DataFrame(self.post_load_process_piece(self._data, unique=True))
+            self.data = pandas.DataFrame(
+                self.post_load_process_piece(pandas.DataFrame(self._data), unique=True)
+            )
             if verbose:
                 print(f"[CubeLogs.load] after postprocessing shape={self.data.shape}")
         elif isinstance(self._data, list) and all(
@@ -614,7 +617,7 @@ class CubeLogs:
 
     def _process_formula(
         self, formula: Union[str, Callable[[pandas.DataFrame], pandas.Series]]
-    ) -> Callable[[pandas.DataFrame], pandas.Series]:
+    ) -> Callable[[pandas.DataFrame], Optional[pandas.Series]]:
         assert callable(formula), f"formula={formula!r} is not supported."
         return formula
 
@@ -625,9 +628,11 @@ class CubeLogs:
         return self.data.shape
 
     @property
-    def columns(self) -> Sequence[str]:
+    def columns(self) -> Sequence[Any]:
         "Returns the columns."
         assert hasattr(self, "data"), "Method load was not called"
+        assert isinstance(self.data, pandas.DataFrame)  # type checking
+        # pyrefly: ignore[bad-return]
         return self.data.columns
 
     def _preprocess(self):
@@ -647,7 +652,7 @@ class CubeLogs:
             )
             assert gr.shape[0] > 0, (
                 f"Something went wrong after the groupby.\n"
-                f"{cp[[*self.keys, self.time, '__index__']].head().T}"
+                f"{cp[[*self.keys_no_time, self.time, '__index__']].head().T}"
             )
             filtered = pandas.merge(cp, gr, on=["__index__", *self.keys_time])
             assert filtered.shape[0] <= self.data.shape[0], (
@@ -797,6 +802,7 @@ class CubeLogs:
             if view_def.agg_multi:
                 append = []
                 for k, f in view_def.agg_multi.items():
+                    # pyrefly: ignore[no-matching-overload]
                     cv = grouped_data.apply(f, include_groups=False)
                     append.append(cv.to_frame(k))
                 data = pandas.concat([data, *append], axis=1)
@@ -1020,8 +1026,10 @@ class CubeLogs:
 
         keys = set(self.keys_no_time) - {columns_to_fix}
         select = data[self.keys_no_time]
+        # pyrefly: ignore[no-matching-overload]
         select_agg = select.groupby(list(keys), as_index=True).apply(
-            lambda x: "-".join(sorted(set(x[columns_to_fix].dropna()))), include_groups=False
+            lambda x: "-".join(sorted(set(x[columns_to_fix].dropna()))),
+            include_groups=False,
         )
         select_agg = select_agg.to_frame(name=columns_to_fix)
         res = pandas.merge(
@@ -1137,6 +1145,7 @@ class CubeLogs:
             if len(nonan) > 0:
                 obs.update(dict(count=len(nonan)))
                 if is_numeric_dtype(nonan) and not pandas.api.types.is_object_dtype(nonan):
+                    # pyrefly: ignore[no-matching-overload]
                     obs.update(
                         dict(
                             min=nonan.min(),
@@ -1208,12 +1217,15 @@ class CubeLogs:
                 df.to_excel(writer, sheet_name=main, freeze_panes=(1, 1))
 
             time_mask_view: Dict[str, pandas.DataFrame] = {}
+            df = None
             for name, view in views.items():
                 if view is None:
                     continue
                 df, tview = self.view(view, return_view_def=True, verbose=max(verbose - 1, 0))
                 if cube_time is not None:
                     cube_mask = cube_time.view(view)
+                    assert isinstance(cube_mask, pandas.DataFrame)  # type checking
+                    assert isinstance(df, pandas.DataFrame)  # type checking
                     aligned = align_dataframe_with(cube_mask, df)
                     if aligned is not None:
                         assert aligned.shape == df.shape, (
@@ -1228,6 +1240,7 @@ class CubeLogs:
                             )
                 if tview is None:
                     continue
+                assert isinstance(df, pandas.DataFrame)  # type checking
                 memory = df.memory_usage(deep=True).sum()
                 if verbose:
                     print(
@@ -1269,7 +1282,9 @@ class CubeLogs:
                         sheet_name=name,
                         freeze_panes=(df.columns.nlevels + 1, df.index.nlevels),
                     )
+                    # pyrefly: ignore[missing-attribute]
                     f_highlights[name] = tview.f_highlight
+                    # pyrefly: ignore[missing-attribute]
                     if tview.plots:
                         plots.append(
                             CubePlot(
@@ -1282,6 +1297,7 @@ class CubeLogs:
                             if self.time in df.columns.names
                             else CubePlot(df, kind="barh", orientation="row", split=True)
                         )
+            assert isinstance(df, pandas.DataFrame)  # type checking
             if raw:
                 assert main not in views, f"{main!r} is duplicated in views {sorted(views)}"
                 # Too long.
@@ -1439,7 +1455,7 @@ class CubeLogs:
             len(configs) >= 2
         ), f"A side by side needs at least two configs but configs={configs}"
         set_keys_time = set(self.keys_time)
-        columns_index = None
+        columns_index: Optional[List[str]] = None
         data_list = []
         for name_conf, conf in configs.items():
             if columns_index is None:
@@ -1478,9 +1494,11 @@ class CubeLogs:
 
         # add metrics
         index_column_name = list(view_res.columns.names).index(column_name)
+        # pyrefly: ignore[missing-attribute]
         index_metrics = list(view_res.columns.names).index("METRICS")
 
         def _mkc(m, s):
+            # pyrefly: ignore[missing-attribute]
             c = ["" for c in view_res.columns.names]
             c[index_column_name] = s
             c[index_metrics] = m
@@ -1515,7 +1533,9 @@ class CubeLogs:
                     ci["CONF"] = iname
                     cj["CONF"] = jname
 
+                    # pyrefly: ignore[bad-index]
                     ci_name = tuple(ci[n] for n in view_res.columns.names)
+                    # pyrefly: ignore[bad-index]
                     cj_name = tuple(cj[n] for n in view_res.columns.names)
                     assert ci_name in view_res.columns or cj_name in view_res.columns, (
                         f"Unable to find column {ci_name} or {cj_name} "
@@ -1562,6 +1582,7 @@ class CubeLogs:
         }
         flat = view_res.groupby(self.time).agg(aggs)
         flat = flat.stack("METRICS", future_stack=True)
+        # pyrefly: ignore[bad-return, missing-attribute]
         return res, flat, view_res.T.sort_index().T
 
 
@@ -1679,7 +1700,7 @@ class CubeLogsPerformance(CubeLogs):
 
     def _process_formula(
         self, formula: Union[str, Callable[[pandas.DataFrame], pandas.Series]]
-    ) -> Callable[[pandas.DataFrame], pandas.Series]:
+    ) -> Callable[[pandas.DataFrame], Optional[pandas.Series]]:
         """
         Processes a formula, converting it into a function.
 
@@ -1726,6 +1747,7 @@ class CubeLogsPerformance(CubeLogs):
                 f"{pprint.pformat(sorted(columns))}"
             )
             # return lambda df: df["time_latency_eager"] / df["time_latency"]
+            # pyrefly: ignore[no-matching-overload]
             return lambda df: pandas.cut(
                 df["speedup"], bins=BUCKET_SCALES, right=False, duplicates="raise"
             )
@@ -1733,9 +1755,9 @@ class CubeLogsPerformance(CubeLogs):
         if formula == "ERR1":
             columns = set(self._filter_column(["^ERR_.*"], self.data.columns))
             if not columns:
-                return lambda df: np.nan
+                return lambda df: None
 
-            def first_err(df: pandas.DataFrame) -> pandas.Series:
+            def first_err(df: pandas.DataFrame) -> Optional[pandas.Series]:
                 ordered = [
                     c
                     for c in [
@@ -1752,7 +1774,7 @@ class CubeLogsPerformance(CubeLogs):
                     ]
                     if c in df.columns
                 ]
-                res = None
+                res: Optional[pandas.Series] = None
                 for c in ordered:
                     if res is None:
                         res = df[c].fillna("")
@@ -1949,6 +1971,7 @@ class CubeLogsPerformance(CubeLogs):
             f"{pprint.pformat(sorted(self.data.columns))}"
         )
 
+    # pyrefly: ignore[bad-override]
     def view(
         self,
         view_def: Optional[Union[str, CubeViewDef]],
@@ -2265,7 +2288,7 @@ class CubeLogsPerformance(CubeLogs):
         if unique:
             return df
         cols = self._filter_column(self._keys, df)
-        res = None
+        res: Optional[pandas.DataFrame] = None
         for c in cols:
             if df[c].isna().any():
                 # Missing values for keys are not supposed to happen.

onnx_diagnostic/helpers/memory_peak.py

@@ -103,6 +103,7 @@ def _process_memory_spy(conn):
     process = psutil.Process(pid)
 
     if cuda:
+        # pyrefly: ignore[missing-import]
         from pynvml import (
             nvmlDeviceGetCount,
             nvmlDeviceGetHandleByIndex,
@@ -131,6 +132,7 @@ def _process_memory_spy(conn):
         mem = process.memory_info().rss
         cpu.update(mem)
         if cuda:
+            # pyrefly: ignore[unbound-name]
             for r, g in zip(gpu_used(), gpus):
                 g.update(r)
         if conn.poll(timeout=timeout):
@@ -142,6 +144,7 @@ def _process_memory_spy(conn):
     end = process.memory_info().rss
     cpu.update(end)
     if cuda:
+        # pyrefly: ignore[unbound-name]
         for r, g in zip(gpu_used(), gpus):
             g.update(r)
 
@@ -151,6 +154,7 @@ def _process_memory_spy(conn):
     for g in gpus:
         g.send(conn)
     if cuda:
+        # pyrefly: ignore[unbound-name]
         nvmlShutdown()
     conn.close()
 
@@ -217,7 +221,7 @@ def start_spying_on(
     Starts the memory spy. The function starts another
     process spying on the one sent as an argument.
 
-    :param pid: process id to spy or the the current one.
+    :param pid: process id to spy or the current one.
     :param delay: delay between two measures.
     :param cuda: True or False to get memory for cuda devices
 

onnx_diagnostic/helpers/mini_onnx_builder.py

@@ -8,11 +8,6 @@ import torch
 from .onnx_helper import dtype_to_tensor_dtype, tensor_dtype_to_np_dtype, from_array_extended
 from . import string_type
 
-STORAGE_TYPE = {
-    TensorProto.FLOAT16: np.int16,
-    TensorProto.BFLOAT16: np.int16,
-}
-
 
 def proto_from_array(
     arr: torch.Tensor,
@@ -67,13 +62,13 @@ def proto_from_array(
         byte_data = (ctypes.c_ubyte * numel * element_size).from_address(np_arr.data_ptr())
         tensor.raw_data = bytes(byte_data)
         if sys.byteorder == "big":
-            np_dtype = tensor_dtype_to_np_dtype(STORAGE_TYPE[tensor.data_type])
-            np.byteswap(np.frombuffer(tensor.raw_data, dtype=np_dtype), inplace=True)
+            np_dtype = tensor_dtype_to_np_dtype(tensor.data_type)
+            np.frombuffer(tensor.raw_data, dtype=np_dtype).byteswap(inplace=True)
     else:
         tensor.raw_data = np_arr.tobytes()
         if sys.byteorder == "big":
             np_dtype = tensor_dtype_to_np_dtype(tensor.data_type)
-            np.byteswap(np.frombuffer(tensor.raw_data, dtype=np_dtype), inplace=True)
+            np.frombuffer(tensor.raw_data, dtype=np_dtype).byteswap(inplace=True)
 
     return tensor
 
@@ -133,6 +128,7 @@ class MiniOnnxBuilder:
                     }
                 shape = tuple(map(int, tensor.shape))
                 self.nodes.append(
+                    # pyrefly: ignore[bad-argument-type]
                     oh.make_node(op_type, [], [name], dtype=dtype, shape=shape, **kwargs)
                 )
                 self.outputs.append(oh.make_tensor_value_info(name, dtype, shape))
@@ -632,6 +628,7 @@ def create_input_tensors_from_onnx_model(
         raise AssertionError(f"Unexpected value for engine={engine!r}")
 
     got = sess.run(None, {})
+    assert isinstance(got, list)  # type checking
     if len(names) == 1:
         name = names[0]
         output = got[0]
@@ -639,12 +636,10 @@ def create_input_tensors_from_onnx_model(
             return None
         if name == "array":
             return output
-        if name == "bool":
-            return bool(output[0])
-        if name == "int":
-            return int(output[0])
-        if name == "float":
-            return float(output[0])
+        if name in {"bool", "int", "float"}:
+            cvt = {"bool": bool, "int": int, "float": float}[name]
+            # pyrefly: ignore[bad-index]
+            return cvt(output[0])
         if name == "tensor":
             return torch.from_numpy(output).to(device)
         assert name.startswith(

onnx_diagnostic/helpers/model_builder_helper.py

@@ -14,7 +14,7 @@ CACHE_SUBDIR = "onnx-diagnostic"
 
 def download_model_builder_to_cache(
     url: str = "https://raw.githubusercontent.com/microsoft/onnxruntime-genai/refs/heads/main/src/python/py/models/builder.py",
-):
+) -> Path:
     """
     Downloads ``builder.py`` from the
     ``https://github.com/microsoft/onnxruntime-genai/blob/main/src/python/py/models/builder.py``.