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

PiperOrigin-RevId: 570055305
google · Oct 2, 2023 · 43db271 · 43db271
1 parent 68bb8e8
commit 43db271
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/praxis/layers/quantization/quantizer.py b/praxis/layers/quantization/quantizer.py
@@ -489,7 +489,7 @@ def _get_optimal_scale_and_min(
       contract_dims: int | Sequence[int],
   ) -> tuple[JTensor, JTensor | None]:
     def quantization_error_and_scale(clipping):
-      q_scale, x_min = self._get_scale_and_min(
+      q_scale, x_min = self._get_scale_and_min(  # pytype: disable=wrong-arg-types  # jnp-type
           x, contract_dims, clipping_coeff=clipping
       )
       x_scaled, zp_time_scale = self._scale(x, q_scale, x_min)