Merge pull request #14131 from jakevdp:mypy-update

PiperOrigin-RevId: 504268346
google · Jan 24, 2023 · bbccf55 · bbccf55
2 parents 7064be1 + 94af71a
commit bbccf55
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -19,7 +19,7 @@ repos:
   - id: mypy
     files: (jax/|tests/typing_test\.py)
     exclude: jax/_src/basearray.py  # Use pyi instead
-    additional_dependencies: [types-requests==2.28.11, jaxlib==0.3.24]
+    additional_dependencies: [types-requests==2.28.11, jaxlib==0.4.1]
 
 - repo: https://github.com/mwouts/jupytext
   rev: v1.14.4

diff --git a/jax/_src/lib/mlir/dialects/__init__.py b/jax/_src/lib/mlir/dialects/__init__.py
@@ -28,4 +28,4 @@
 if use_stablehlo:
   import jaxlib.mlir.dialects.stablehlo as hlo
 else:
-  import jaxlib.mlir.dialects.mhlo as hlo
+  import jaxlib.mlir.dialects.mhlo as hlo  # type: ignore[no-redef]
diff --git a/jax/experimental/rnn.py b/jax/experimental/rnn.py
@@ -97,7 +97,7 @@
 try:
   from jax._src.lib import gpu_rnn
 except ImportError:
-  gpu_rnn = None
+  gpu_rnn = None  # type: ignore[assignment]
 
 PRNGKeyArray = Any
 sigmoid = jax.nn.sigmoid