minor modifications (#392)

* minor modifications

* fix

* fix

* types

* fix

* miss

* fix

* fix

* fix doc

Author: xadupre

.github/workflows/pyrefly.yaml

@@ -21,6 +21,7 @@ jobs:
         run: |
           pip install pyrefly
           pip install -r requirements.txt
+          pip install transformers pandas matplotlib openpyxl
 
       - name: Run pyrefly
         run: pyrefly check

_doc/conf.py

@@ -132,6 +132,7 @@ def linkcode_resolve(domain, info):
     ("py:class", "Module"),
     ("py:class", "np.ndarray"),
     ("py:class", "onnx_ir.Tuple"),
+    ("py:class", "pandas.api.typing.DataFrameGroupBy"),
     ("py:class", "pandas.core.groupby.generic.DataFrameGroupBy"),
     ("py:class", "pipeline.Pipeline"),
     ("py:class", "torch._guards.Source"),

_doc/recipes/plot_dynamic_shapes_json.py

@@ -81,7 +81,7 @@ def flatten_unflatten_like_dynamic_shapes(obj):
         subtrees.append(value)
         start = end
     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)

onnx_diagnostic/export/validate.py

@@ -80,6 +80,7 @@ def _get(a):
         )
     got = modep(*_get(args), **_get(kwargs))
     if verbose:
+        # pyrefly: ignore[unbound-name]
         d = time.perf_counter() - begin
         print(f"[compare_modules] done in {d} with output={string_type(got, with_shape=True)}")
     if mod:
@@ -89,6 +90,7 @@ def _get(a):
         expected = mod(*_get(args), **_get(kwargs))
         diff = max_diff(expected, got)
         if verbose:
+            # pyrefly: ignore[unbound-name]
             d = time.perf_counter() - begin
             print(
                 f"[compare_modules] done in {d} with "

onnx_diagnostic/ext_test_case.py

@@ -780,7 +780,7 @@ def unit_test_going(self) -> bool:
 
     @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

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/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)

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 @@ def __init__(
             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 @@ def _make_model(self, itype: int, rank: int, has_bias: bool) -> onnx.ModelProto:
             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 @@ def __init__(
         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 @@ def _make_model(self, itype: int, ranka: int, rankb: int) -> onnx.ModelProto:
             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 @@ def run(self, a, b):
         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 @@ def computation_order(
         :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 @@ def text_positions(
     @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,15 +151,15 @@ def start_names(self) -> List[onnx.NodeProto]:
             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]
 
     @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 @@ def input_names(self) -> List[str]:
     @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 [