jax.numpy.clip: update use of deprecated arguments.

- a is now positional-only
- a_min is now min
- a_max is now max

The old argument names have been deprecated since JAX v0.4.27.

PiperOrigin-RevId: 714108629

baselines/t5/decoding.py

@@ -846,7 +846,7 @@ def beam_search_loop_body_fn(state: BeamState) -> BeamState:
   # Clip the last value caused by the `diff` operator of the padding
   # value 0.0 and the last cumulative sum. That value is positive because
   # the sum of all log probabilities is negative.
-  finished_token_scores = jnp.clip(finished_token_scores, a_max=0.)
+  finished_token_scores = jnp.clip(finished_token_scores, max=0.)
 
   if return_token_scores:
     return finished_seqs[:, :, 1:], finished_token_scores