Another batch of Pyrefly fixes

Author: superbobry

jax/_src/export/_export.py

@@ -357,7 +357,7 @@ def deserialize(blob: bytearray) -> Exported:
 
 
 class _SerializeAuxData(Protocol):
-  def __call__(self, aux_data: PyTreeAuxData) -> bytes:
+  def __call__(self, aux_data: PyTreeAuxData, /) -> bytes:
     """Serializes the PyTree node AuxData.
 
     The AuxData is returned by the ``flatten_func`` registered by
@@ -366,7 +366,7 @@ def __call__(self, aux_data: PyTreeAuxData) -> bytes:
 
 
 class _DeserializeAuxData(Protocol):
-  def __call__(self, serialized_aux_data: bytes) -> PyTreeAuxData:
+  def __call__(self, serialized_aux_data: bytes, /) -> PyTreeAuxData:
     """Deserializes the PyTree node AuxData.
 
     The result will be passed to ``_BuildFromChildren``.
@@ -497,6 +497,7 @@ def register_namedtuple_serialization(
   def serialize_auxdata(aux_data: PyTreeAuxData) -> bytes:
     # Store the serialized keys in the serialized auxdata
     del aux_data
+    # pyrefly: ignore[missing-attribute]
     return json.dumps(nodetype._fields).encode("utf-8")
 
   def deserialize_auxdata(serialized_aux_data: bytes) -> PyTreeAuxData:
@@ -901,8 +902,8 @@ def _module_to_bytecode(module: ir.Module) -> bytes:
   # Note that this does not verify any JAX custom calls, which are only
   # guaranteed 3w of forward compatibility, and only prevents use of new
   # StableHLO features from failing on older hardware.
-  target_version = hlo.get_version_from_compatibility_requirement(
-    hlo.StablehloCompatibilityRequirement.WEEK_4)
+  target_version = hlo.get_version_from_compatibility_requirement(  # pyrefly: ignore[missing-attribute]
+    hlo.StablehloCompatibilityRequirement.WEEK_4)  # pyrefly: ignore[missing-attribute]
   module_serialized = xla_client._xla.mlir.serialize_portable_artifact(  # type: ignore
       mlir_str, target_version, xb.get_backend().serialize_with_sdy)
   return module_serialized
@@ -995,7 +996,7 @@ def is_token(typ, attrs):
       new_arg_attrs = []
       for idx in new_main_arg_indices:
         new_arg_attr = {}
-        for attr in arg_attrs[idx]:
+        for attr in arg_attrs[idx]:  # pyrefly: ignore[not-iterable]
           if attr.name == "tf.aliasing_output":
             i = new_main_result_indices.index(attr.attr.value)
             new_arg_attr[attr.name] = ir.IntegerAttr.get(
@@ -1374,7 +1375,7 @@ def flattened_primal_fun_jax(*args_flat):
     if has_named_shardings or mesh:
       vjp_in_shardings = tuple(
           _get_named_sharding(has_named_shardings, named_sharding,  # type: ignore
-                              hlo_sharding, aval, mesh)
+                              hlo_sharding, aval, mesh)  # pyrefly: ignore[bad-argument-type]
           for named_sharding, hlo_sharding, aval in zip(
             itertools.chain(in_named_shardings, out_named_shardings),
             itertools.chain(in_shardings_hlo, out_shardings_hlo),
@@ -1548,8 +1549,8 @@ def _call_exported_impl(*args, exported: Exported):
 def get_mesh_from_symbol(symtab: ir.SymbolTable) -> mesh_lib.AbstractMesh:
   if "mesh" not in symtab:
     return mesh_lib.empty_abstract_mesh
-  mesh_attr = sdy.MeshAttr(symtab["mesh"].mesh)
-  axes = [sdy.MeshAxisAttr(a) for a in mesh_attr.axes]
+  mesh_attr = sdy.MeshAttr(symtab["mesh"].mesh)  # pyrefly: ignore[missing-attribute]
+  axes = [sdy.MeshAxisAttr(a) for a in mesh_attr.axes]  # pyrefly: ignore[missing-attribute]
   if not axes:
     return mesh_lib.empty_abstract_mesh
   axes_sizes = tuple(a.size for a in axes)

jax/_src/export/serialization.py

@@ -17,10 +17,10 @@
 from __future__ import annotations
 
 import types
-from collections.abc import Callable, Sequence
+from collections.abc import Callable, Iterable
 import itertools
 from functools import partial
-from typing import Any, TypeVar
+from typing import cast, Any, TypeVar
 
 try:
   import flatbuffers
@@ -166,9 +166,10 @@ def _serialize_exported(
 def _serialize_array(
     builder: flatbuffers.Builder,
     serialize_one: Callable[[flatbuffers.Builder, T], int],
-    elements: Sequence[T],
+    elements: Iterable[T],
 ) -> int:
   element_offsets = [serialize_one(builder, e) for e in elements]
+  del elements
   ser_flatbuf.PyTreeDefStartChildrenVector(builder, len(element_offsets))
   for sc in reversed(element_offsets):
     builder.PrependUOffsetTRelative(sc)
@@ -216,8 +217,10 @@ def _deserialize_exported(exp: ser_flatbuf.Exported) -> _export.Exported:
                                   partial(_deserialize_aval, scope=scope, sharding=None))
     out_avals = _deserialize_tuple(exp.OutAvalsLength, exp.OutAvals,
                                    partial(_deserialize_aval, scope=scope, sharding=None))
-    in_shardings_hlo, in_shardings = in_shardings, (None,) * len(in_shardings)  # type: ignore
-    out_shardings_hlo, out_shardings = out_shardings, (None,) * len(out_shardings)  # type: ignore
+    in_shardings_hlo = cast(tuple[_export.HloSharding | None, ...], in_shardings)
+    in_shardings = (None,) * len(in_shardings)  # type: ignore
+    out_shardings_hlo = cast(tuple[_export.HloSharding | None, ...], out_shardings)
+    out_shardings = (None,) * len(out_shardings)  # type: ignore
   platforms = _deserialize_tuple(
       exp.PlatformsLength,
       exp.Platforms,
@@ -313,6 +316,7 @@ def serialize_key(builder, k):
         builder, serialize_key, node_data[1]
     )
   elif node_type in _export.serialization_registry:
+    assert node_type is not None
     kind = ser_flatbuf.PyTreeDefKind.custom
     serialized_name, serialize_auxdata = _export.serialization_registry[node_type]
     custom_name = builder.CreateString(serialized_name)
@@ -497,7 +501,7 @@ def _deserialize_partition_spec_one_axis(
 def _serialize_partition_spec(builder: flatbuffers.Builder,
                               spec: partition_spec.PartitionSpec) -> int:
   partitions = _serialize_array(builder, _serialize_partition_spec_one_axis,
-                                spec._partitions)
+                                spec._partitions)  # pyrefly: ignore[bad-argument-type]
   reduced = _serialize_array(builder,  # type: ignore
                              lambda builder, ps: builder.CreateString(ps),
                              spec.reduced)
@@ -583,8 +587,9 @@ def _deserialize_aval(aval: ser_flatbuf.AbstractValue, *,
   else:
     mem_space = core.MemorySpace.Device
 
-  aval = core.ShapedArray(shape, dtype, memory_space=mem_space)
-  return core.update_aval_with_sharding(aval, sharding)
+  return core.update_aval_with_sharding(
+      core.ShapedArray(shape, dtype, memory_space=mem_space), sharding
+  )
 
 
 def _serialize_sharding(

jax/_src/export/serialization_generated.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# pyrefly: ignore-errors
 # pytype: skip-file
 # automatically generated by the FlatBuffers compiler, do not modify
 

jax/_src/export/shape_poly.py

@@ -28,7 +28,7 @@
 import copy
 import operator as op
 import tokenize
-from typing import Any, Union, overload
+from typing import Any, TypeAlias, Union, overload
 import warnings
 
 import numpy as np
@@ -52,8 +52,8 @@
 DType = Any
 
 # Tuples of terms and their coefficients, sorted with the largest term first.
-SortedTerms = Sequence[tuple["_DimTerm", int]]
-SortedFactors = Sequence[tuple["_DimFactor", int]]
+SortedTerms: TypeAlias = Sequence[tuple["_DimTerm", int]]
+SortedFactors: TypeAlias = Sequence[tuple["_DimFactor", int]]
 
 # Normalization rules represent the explicit constraint `t*tk == e` as
 # a mapping of `t` to `(e, tk)`.
@@ -583,14 +583,14 @@ def _get_vars(self) -> set[str]:
   @overload
   @staticmethod
   def _linear_combination_sorted_pairs(
-      e1: SortedTerms, i1: int, f1: int,
-      e2: SortedTerms, i2: int, f2: int) -> SortedTerms:  ...  # type: ignore[bad-return-type,unused-ignore]
+      pairs1: SortedTerms, i1: int, f1: int,
+      pairs2: SortedTerms, i2: int, f2: int) -> SortedTerms:  ...  # type: ignore[bad-return-type,unused-ignore]
 
   @overload
   @staticmethod
   def _linear_combination_sorted_pairs(
-      e1: SortedFactors, i1: int, f1: int,
-      e2: SortedFactors, i2: int, f2: int) -> SortedFactors:  ...  # type: ignore[bad-return-type,unused-ignore]
+      pairs1: SortedFactors, i1: int, f1: int,
+      pairs2: SortedFactors, i2: int, f2: int) -> SortedFactors:  ...  # type: ignore[bad-return-type,unused-ignore]
 
   @staticmethod
   def _linear_combination_sorted_pairs(
@@ -862,7 +862,7 @@ def __gt__(self, other: DimSize):
   def __lt__(self, other: DimSize):
     return not _geq_decision(self, other, lambda: f"'{self}' < '{other}'")
 
-  def _divmod(self, divisor: DimSize) -> tuple[DimSize, int]:
+  def _divmod(self, divisor: DimSize) -> tuple[DimSize, DimSize]:
     """
     Floor division with remainder (divmod) generalized to expressions.
     If the `divisor` is not a constant, the remainder must be 0.
@@ -1626,7 +1626,7 @@ def expr(self, tok: tokenize.TokenInfo) -> tuple[DimSize, tokenize.TokenInfo]:
       tok = self.next_tok()
     elif tok.exact_type == tokenize.PLUS:
       tok = self.next_tok()
-    acc = None
+    acc: DimSize | None = None
     while True:
       t, tok = self.term(tok)
       t_sign = - t if next_t_negated else t
@@ -1640,7 +1640,7 @@ def expr(self, tok: tokenize.TokenInfo) -> tuple[DimSize, tokenize.TokenInfo]:
   FOLLOW_TERM = FOLLOW_EXPR + [tokenize.PLUS, tokenize.MINUS]
   def term(self, tok: tokenize.TokenInfo) -> tuple[DimSize, tokenize.TokenInfo]:
     # A term is product of factors. Each factor may be raised to an integer power.
-    acc = None
+    acc: DimSize | None = None
     while True:
       f, tok = self.factor(tok)
       if tok.exact_type == tokenize.CIRCUMFLEX:

jax/_src/internal_test_util/export_back_compat_test_data/cpu_triangular_solve_blas_trsm.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# pyrefly: ignore-errors
 # ruff: noqa
 
 import datetime

jax/_src/internal_test_util/export_back_compat_test_data/cpu_tridiagonal_solve_lapack_gtsv.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# pyrefly: ignore-errors
 # ruff: noqa
 
 import datetime

jax/_src/internal_test_util/export_back_compat_test_data/cuda_cholesky_solver_potrf.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# pyrefly: ignore-errors
 # ruff: noqa
 
 import datetime

jax/_src/internal_test_util/export_back_compat_test_data/cuda_eigh_cusolver_syev.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# pyrefly: ignore-errors
 # ruff: noqa
 
 import datetime

jax/_src/internal_test_util/export_back_compat_test_data/cuda_lu_cusolver_getrf.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# pyrefly: ignore-errors
 # ruff: noqa
 
 import datetime

jax/_src/internal_test_util/export_back_compat_test_data/cuda_qr_cusolver_geqrf.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# pyrefly: ignore-errors
 # ruff: noqa
 
 import datetime