Merge pull request jax-ml#24420 from superbobry:maint-2

PiperOrigin-RevId: 688271404

Author: Google-ML-Automation

.pre-commit-config.yaml

@@ -26,12 +26,12 @@ repos:
     files: \.py$
 
 - repo: https://github.com/astral-sh/ruff-pre-commit
-  rev: 8b5112a3b2ad121439a2092f8ff548c0d80f2514  # frozen: v0.6.1
+  rev: 8983acb92ee4b01924893632cf90af926fa608f0  # frozen: v0.7.0
   hooks:
   - id: ruff
 
 - repo: https://github.com/pre-commit/mirrors-mypy
-  rev: 'd4911cfb7f1010759fde68da196036feeb25b99d'  # frozen: v1.11.2
+  rev: '102bbee94061ff02fd361ec29c27b7cb26582f5f'  # frozen: v1.12.2
   hooks:
   - id: mypy
     files: (jax/|tests/typing_test\.py)

jax/_src/core.py

@@ -3270,7 +3270,7 @@ def pp_eqn(eqn: JaxprEqn, context: JaxprPpContext, settings: JaxprPpSettings
            ) -> pp.Doc:
   rule = (_pp_eqn if not settings.custom_pp_eqn_rules else
           pp_eqn_rules.get(eqn.primitive, _pp_eqn))
-  doc = rule(eqn, context, settings)  # type: ignore[operator]
+  doc = rule(eqn, context, settings)
   user_frame = source_info_util.user_frame(eqn.source_info)
   return doc if user_frame is None else pp.source_map(doc, user_frame)
 

jax/_src/interpreters/partial_eval.py

@@ -2738,7 +2738,7 @@ def inline_jaxpr_into_trace(
     outvars = [Var('', v.aval) for v in eqn.outvars]
     src_ = (src if not eqn.source_info.name_stack else
             src.replace(name_stack=src.name_stack + eqn.source_info.name_stack))
-    trace.frame.add_eqn(eqn.replace(invars, outvars, source_info=src_))  # type: ignore
+    trace.frame.add_eqn(eqn.replace(invars, outvars, source_info=src_))
     map(env.setdefault, eqn.outvars, outvars)
 
   tracer_env: dict[Var, Any] = dict(zip([*jaxpr.constvars, *jaxpr.invars],

jax/_src/interpreters/pxla.py

@@ -1750,7 +1750,7 @@ def _get_and_check_device_assignment(
   elif first_sharding_info is None:
     final_device_assignment = (_get_default_device(),)
   else:
-    final_device_assignment = first_sharding_info[0]
+    final_device_assignment = first_sharding_info[0]  # type: ignore
   return xb.get_device_backend(final_device_assignment[0]), final_device_assignment
 
 MaybeSharding = Union[JSharding, UnspecifiedValue]

jax/_src/numpy/lax_numpy.py

@@ -6477,7 +6477,7 @@ def _arange(start: ArrayLike | DimSize, stop: ArrayLike | DimSize | None = None,
     if (not dtypes.issubdtype(start_dtype, np.integer) and
         not dtypes.issubdtype(start_dtype, dtypes.extended)):
       ceil_ = ufuncs.ceil if isinstance(start, core.Tracer) else np.ceil
-      start = ceil_(start).astype(int)  # type: ignore[operator]
+      start = ceil_(start).astype(int)
     return lax.iota(dtype, start)  # type: ignore[arg-type]
   else:
     if step is None and start == 0 and stop is not None:

jax/_src/pallas/pallas_call.py

@@ -254,7 +254,7 @@ def _pallas_call_impl_interpret(
   num_inout_blocks = len(block_args) + len(out)
   grid_start_indices = (jnp.int32(0),) * len(grid)
   if grid:
-    num_iterations = reduce(jnp.multiply, grid)
+    num_iterations = reduce(jnp.multiply, grid)  # type: ignore[arg-type]
   else:
     # Base case is always one iteration when grid is ()
     num_iterations = 1
@@ -1174,7 +1174,7 @@ def pallas_call_checkify_oob_grid(error: checkify.Error,
   )
   grid_start_indices = (jnp.int32(0),) * len(grid)
   if grid:
-    num_iterations = reduce(jnp.multiply, grid)
+    num_iterations = reduce(jnp.multiply, grid)  # type: ignore[arg-type]
   else:
     # Base case is always one iteration when grid is ()
     num_iterations = 1

jax/_src/pjit.py

@@ -779,7 +779,7 @@ def _extract_implicit_args(
           args[d1.val] = d2
         assert core.same_referent(args[d1.val], d2)
   assert all(x is not None for x in args)
-  return [x for x, (_, e) in zip(args, in_type) if not e]  # pytype: disable=bad-return-type
+  return [x for x, (_, e) in zip(args, in_type) if not e]  # type: ignore
 
 def _flat_axes_specs(abstracted_axes, *args, **kwargs
                      ) -> list[pe.AbstractedAxesSpec] | None:
@@ -1545,6 +1545,7 @@ def _resolve_in_shardings(args, pjit_in_shardings: Sequence[PjitSharding]
           else:
             resolved_in_shardings.append(arg_s)
         else:
+          assert isinstance(arg_s, sharding.Sharding)
           if dispatch.is_single_device_sharding(arg_s):
             resolved_in_shardings.append(UNSPECIFIED)
           else:
@@ -1581,7 +1582,7 @@ def _resolve_in_shardings(args, pjit_in_shardings: Sequence[PjitSharding]
             not isinstance(arg_s, PmapSharding) and
             not op_shardings.are_op_shardings_equal(
                 pjit_in_s._to_xla_hlo_sharding(arg.ndim),  # type: ignore
-                arg_s._to_xla_hlo_sharding(arg.ndim))):
+                arg_s._to_xla_hlo_sharding(arg.ndim))):  # type: ignore
           raise ValueError('Sharding passed to pjit does not match the sharding '
                            'on the respective arg. '
                            f'Got pjit sharding: {pjit_in_s},\n'