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

PiperOrigin-RevId: 568283267

Author: Jake VanderPlas

kfac_jax/_src/optimizer.py

@@ -865,7 +865,7 @@ def _coefficients_and_quad_change(
       else:
         quad_change = jnp.nan
 
-      return coefficients, quad_change
+      return coefficients, quad_change  # pytype: disable=bad-return-type  # jnp-type
 
   @utils.auto_scope_method
   def _update_damping(

kfac_jax/_src/patches_second_moment.py

@@ -628,7 +628,7 @@ def general_loop_body(i, image):
       else:
         wf_n = weighting_array[in_spec.n_axis]
         wf_spatial = [weighting_array.shape[a] for a in in_spec.spatial_axes]
-        wf_sizes = in_spec.create_shape(wf_n, jnp.ones([]), *wf_spatial)
+        wf_sizes = in_spec.create_shape(wf_n, jnp.ones([]), *wf_spatial)  # pytype: disable=wrong-arg-types  # jnp-type
         wf_i = _slice_array(weighting_array, index, wf_sizes)
     else:
       wf_i = None

kfac_jax/_src/utils/types.py

@@ -81,4 +81,4 @@ def get_float_dtype_and_check_consistency(obj: ArrayTree) -> DType:
     else:
       raise ValueError("Non-float dtype detected.")
 
-  return dtype
+  return dtype  # pytype: disable=bad-return-type  # jnp-type