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

PiperOrigin-RevId: 568475283 Change-Id: Ice9d8b610d2d8ab1c541679590e184ab91cbc05e
google-deepmind · Nov 24, 2023 · d592df1 · d592df1
1 parent 1aa0a2a
commit d592df1
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/ferminet/jastrows.py b/ferminet/jastrows.py
@@ -41,8 +41,8 @@ def _jastrow_ee(
       for r in jnp.split(r_ee, nspins[0:1], axis=0)
   ]
   r_ees_parallel = jnp.concatenate([
-      r_ees[0][0][jnp.triu_indices(nspins[0], k=1)],
-      r_ees[1][1][jnp.triu_indices(nspins[1], k=1)],
+      r_ees[0][0][jnp.triu_indices(nspins[0], k=1)],  # pytype: disable=wrong-arg-types  # jnp-type
+      r_ees[1][1][jnp.triu_indices(nspins[1], k=1)],  # pytype: disable=wrong-arg-types  # jnp-type
   ])
 
   if r_ees_parallel.shape[0] > 0: