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

PiperOrigin-RevId: 568300084

rax/_src/segment_utils.py

@@ -67,7 +67,7 @@ def segment_log_softmax(
     a: Array, segments: Array, where: Optional[Array] = None
 ) -> Array:
   """Returns segment log-softmax."""
-  a_max = segment_max(a, segments, where=where, initial=jnp.min(a))
+  a_max = segment_max(a, segments, where=where, initial=jnp.min(a))  # pytype: disable=wrong-arg-types  # jnp-type
   shifted = a - jax.lax.stop_gradient(a_max)
   shifted_logsumexp = jnp.log(
       segment_sum(jnp.exp(shifted), segments, where=where)
@@ -79,7 +79,7 @@ def segment_softmax(
     a: Array, segments: Array, where: Optional[Array] = None
 ) -> Array:
   """Returns segment softmax."""
-  a_max = segment_max(a, segments, where=where, initial=jnp.min(a))
+  a_max = segment_max(a, segments, where=where, initial=jnp.min(a))  # pytype: disable=wrong-arg-types  # jnp-type
   unnormalized = jnp.exp(a - jax.lax.stop_gradient(a_max))
   return unnormalized / segment_sum(unnormalized, segments, where=where)