[JAX] Replace jnp.clip(..., a_min=..., a_max=...) with jnp.clip(..., min=..., max=...).

a_min and a_max are deprecated parameter names to jax.numpy.clip.

PiperOrigin-RevId: 890539968

Author: hawkinsp

src/maxtext/layers/moe.py

@@ -598,8 +598,14 @@ def apply_ffn_activation(self, layer_w0, layer_w1):
     """Applies FFN activation function."""
     with jax.named_scope("ffn_act"):
       if self.config.decoder_block == ctypes.DecoderBlockType.GPT_OSS:
-        layer_w0 = jnp.clip(layer_w0, a_min=None, a_max=self.config.mlp_activations_limit)
-        layer_w1 = jnp.clip(layer_w1, a_min=-self.config.mlp_activations_limit, a_max=self.config.mlp_activations_limit)
+        layer_w0 = jnp.clip(
+            layer_w0, min=None, max=self.config.mlp_activations_limit
+        )
+        layer_w1 = jnp.clip(
+            layer_w1,
+            min=-self.config.mlp_activations_limit,
+            max=self.config.mlp_activations_limit,
+        )
         layer_act = self.activation_fn(layer_w0 * 1.702)
         glu = jnp.multiply(layer_w0, layer_act)
         intermediate_layer = jnp.multiply(glu, (layer_w1 + 1))

tests/utils/forward_pass_logit_checker.py

@@ -332,8 +332,14 @@ def main(config, test_args):  # pylint: disable=W0621
       max_logging.log(msg)
 
       if test_args.clip_logits_epsilon is not None:
-        model_probabilities = jnp.clip(jax.nn.softmax(train_logits_slice, axis=-1), a_min=test_args.clip_logits_epsilon)
-        golden_probabilities = jnp.clip(jax.nn.softmax(golden_logits_slice, axis=-1), a_min=test_args.clip_logits_epsilon)
+        model_probabilities = jnp.clip(
+            jax.nn.softmax(train_logits_slice, axis=-1),
+            min=test_args.clip_logits_epsilon,
+        )
+        golden_probabilities = jnp.clip(
+            jax.nn.softmax(golden_logits_slice, axis=-1),
+            min=test_args.clip_logits_epsilon,
+        )
       else:
         model_probabilities = jax.nn.softmax(train_logits_slice, axis=-1)
         golden_probabilities = jax.nn.softmax(golden_logits_slice, axis=-1)