[LSC] Ignore incorrect type annotations related to jax.numpy APIs

PiperOrigin-RevId: 568576099
google · Sep 26, 2023 · 593bbf0 · 593bbf0
1 parent 675fa2a
commit 593bbf0
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/jax_cfd/base/array_utils.py b/jax_cfd/base/array_utils.py
@@ -344,7 +344,7 @@ def orthogonalize(vector, others):
   return q, r
 
 
-def interp1d(
+def interp1d(  # pytype: disable=annotation-type-mismatch  # jnp-type
     x: Array,
     y: Array,
     axis: int = -1,

diff --git a/jax_cfd/spectral/utils.py b/jax_cfd/spectral/utils.py
@@ -178,7 +178,7 @@ def vorticity_to_velocity(
   kx, ky = grid.rfft_mesh()
   two_pi_i = 2 * jnp.pi * 1j
   laplace = two_pi_i ** 2 * (abs(kx)**2 + abs(ky)**2)
-  laplace = laplace.at[0, 0].set(1)
+  laplace = laplace.at[0, 0].set(1)  # pytype: disable=attribute-error  # jnp-type
 
   def ret(vorticity_hat):
     psi_hat = -1 / laplace * vorticity_hat