Fix pre-commit to run on all files in CI.

Author: dfm

.github/workflows/ci-build.yaml

@@ -39,7 +39,7 @@ jobs:
         with:
           path: ~/.cache/pre-commit
           key: pre-commit-${{ env.pythonLocation }}-${{ hashFiles('.pre-commit-config.yaml', 'setup.py') }}
-      - run: pre-commit run --show-diff-on-failure --color=always
+      - run: pre-commit run --show-diff-on-failure --color=always --all-files
 
   build:
     name: "build ${{ matrix.name-prefix }} (py ${{ matrix.python-version }} on ubuntu-20.04, x64=${{ matrix.enable-x64}})"

jax/_src/array.py

@@ -1173,7 +1173,7 @@ def _array_shard_arg(xs, shardings, layouts, copy_semantics):
     copy_outs = xc.batched_copy_array_to_devices_with_sharding(
         batch_xs, batch_devs, batch_shardings, batch_cs)
   else:
-    copy_outs = xc.batched_copy_array_to_devices_with_sharding(  # type: ignore
+    copy_outs = xc.batched_copy_array_to_devices_with_sharding(  # pytype: disable=missing-parameter
         batch_xs, batch_devs, batch_shardings)
   for i, copy_out in safe_zip(batch_indices, copy_outs):
     assert results[i] is None

jax/_src/config.py

@@ -22,7 +22,7 @@
 import os
 import sys
 import threading
-from typing import Any, Generic, NamedTuple, NoReturn, Optional, Protocol, TypeVar, cast, TYPE_CHECKING
+from typing import Any, Generic, NamedTuple, NoReturn, Optional, Protocol, TypeVar, cast
 
 from jax._src import lib
 from jax._src.lib import guard_lib
@@ -371,7 +371,7 @@ class _Unset: pass
 
   _thread_local_state = threading.local()
 
-  class State(Generic[_T]):
+  class State(Generic[_T]):  # type: ignore[no-redef]
 
     __slots__ = (
         '_name', '_value', '_update_thread_local_hook', '_update_global_hook',

jax/experimental/mosaic/gpu/dialect_lowering.py

@@ -29,6 +29,8 @@
 from jaxlib.mlir.dialects import nvvm
 from .utils import c, memref_ptr, single_thread_predicate
 
+# mypy: ignore-errors
+
 
 MlirLoweringRule = Callable[[ir.Operation | ir.OpView], Sequence[ir.Value]]
 

jax/experimental/mosaic/gpu/fragmented_array.py

@@ -19,6 +19,7 @@
 import dataclasses
 import functools
 import math
+from collections.abc import Callable
 from typing import Iterable, Sequence, TypeVar
 
 import jax