Repository: google/jax

Issue:
pre-commit mypy check seems outdated
https://github.com/google/jax/blob/375777f43c919dd4133c4b438c31f920e0dad6a6/.pre-commit-config.yaml#L16-L21

It seems that it use `jaxlib==0.1.74`.
In addition, `mypy` version can be updated to `v0.942`.


Discussion:
And I can't pass mypy typecheck locally on main branch
```
jax/_src/lib/__init__.py:111: error: Skipping analyzing "jaxlib.cpu_feature_guard": module is installed, but missing library stubs or py.typed marker  [import]
jax/_src/lib/__init__.py:111: note: See https://mypy.readthedocs.io/en/stable/running_mypy.html#missing-imports
jax/_src/lib/__init__.py:124: error: Skipping analyzing "jaxlib.cusolver": module is installed, but missing library stubs or py.typed marker  [import]
jax/_src/lib/__init__.py:129: error: Skipping analyzing "jaxlib.hipsolver": module is installed, but missing library stubs or py.typed marker  [import]
jax/_src/lib/__init__.py:134: error: Skipping analyzing "jaxlib.cusparse": module is installed, but missing library stubs or py.typed marker  [import]
jax/_src/lib/__init__.py:139: error: Skipping analyzing "jaxlib.hipsparse": module is installed, but missing library stubs or py.typed marker  [import]
jax/_src/lib/__init__.py:147: error: Skipping analyzing "jaxlib.cuda_prng": module is installed, but missing library stubs or py.typed marker  [import]
jax/_src/lib/__init__.py:152: error: Skipping analyzing "jaxlib.hip_prng": module is installed, but missing library stubs or py.typed marker  [import]
jax/_src/lib/__init__.py:157: error: Skipping analyzing "jaxlib.cuda_linalg": module is installed, but missing library stubs or py.typed marker  [import]
jax/_src/lib/__init__.py:162: error: Skipping analyzing "jaxlib.hip_linalg": module is installed, but missing library stubs or py.typed marker  [import]
jax/interpreters/xla.py:160: error: Argument 1 to "list" has incompatible type "Union[Tuple[int, ...], Tuple[Optional[Tuple[int, ...]], ...]]"; expected "Iterable[int]"  [arg-type]
jax/_src/nn/functions.py:382: error: Argument 1 to "index" has incompatible type "Union[int, Hashable]"; expected "SupportsIndex"  [arg-type]
jax/experimental/jax2tf/converters_eval/examples.py:19: error: Cannot find implementation or library stub for module named "flax"  [import]
jax/experimental/gda_serialization/serialization.py:26: error: Cannot find implementation or library stub for module named "tensorstore"  [import]
jax/tools/jax_to_ir.py:86: error: Cannot find implementation or library stub for module named "tensorflow"  [import]
jax/experimental/jax2tf/converters_eval/converters.py:25: error: Cannot find implementation or library stub for module named "tensorflow"  [import]
jax/experimental/jax2tf/converters_eval/converters.py:26: error: Cannot find implementation or library stub for module named "tensorflowjs.converters"  [import]
Found 16 errors in 7 files (checked 257 source files)
```

Resolving patch:
diff --git a/jax/_src/nn/functions.py b/jax/_src/nn/functions.py
--- a/jax/_src/nn/functions.py
+++ b/jax/_src/nn/functions.py
@@ -379,7 +379,7 @@ def _one_hot(x: Array, num_classes: int, *,
                        f"but {num_classes} != {axis_size}") from None
     axis_idx = lax.axis_index(axis)
     return jnp.asarray(x == axis_idx, dtype=dtype)
-  axis = operator.index(axis)
+  axis = operator.index(axis)  # type: ignore[arg-type]
   lhs = lax.expand_dims(x, (axis,))
   rhs_shape = [1] * x.ndim
   rhs_shape.insert(output_pos_axis, num_classes)
diff --git a/jax/interpreters/xla.py b/jax/interpreters/xla.py
--- a/jax/interpreters/xla.py
+++ b/jax/interpreters/xla.py
@@ -157,7 +157,7 @@ def sharding_to_proto(sharding: SpatialSharding):
     proto.type = xc.OpSharding.Type.REPLICATED
   else:
     proto.type = xc.OpSharding.Type.OTHER
-    proto.tile_assignment_dimensions = list(sharding)
+    proto.tile_assignment_dimensions = list(sharding)  # type: ignore
     proto.tile_assignment_devices = list(range(np.product(sharding)))  # type: ignore
   return proto
 


Test patch: