update mypy & related package versions

MarMarhoun · Apr 15, 2022 · 7972b98 · 7972b98
1 parent 375777f
commit 7972b98
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 4 deletions.
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -14,11 +14,11 @@ repos:
   - id: flake8
 
 - repo: https://github.com/pre-commit/mirrors-mypy
-  rev: 'v0.931'
+  rev: 'v0.942'
   hooks:
   - id: mypy
     files: jax/
-    additional_dependencies: [types-requests==0.1.11, jaxlib==0.1.74]
+    additional_dependencies: [types-requests==2.27.16, jaxlib==0.3.5]
 
 - repo: https://github.com/mwouts/jupytext
   rev: v1.13.6

diff --git 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
@@ -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