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: 715343439
Change-Id: I50b086b249360c142f42ed4d8e50d48692c11dfa

alphafold/model/modules.py

@@ -1854,8 +1854,8 @@ def __call__(self, batch, is_training, safe_key=None):
       rel_pos = jax.nn.one_hot(
           jnp.clip(
               offset + c.max_relative_feature,
-              a_min=0,
-              a_max=2 * c.max_relative_feature),
+              min=0,
+              max=2 * c.max_relative_feature),
           2 * c.max_relative_feature + 1)
       pair_activations += common_modules.Linear(
           c.pair_channel, name='pair_activiations')(

alphafold/model/modules_multimer.py

@@ -564,7 +564,7 @@ def _relative_encoding(self, batch):
     dtype = jnp.bfloat16 if gc.bfloat16 else jnp.float32
 
     clipped_offset = jnp.clip(
-        offset + c.max_relative_idx, a_min=0, a_max=2 * c.max_relative_idx)
+        offset + c.max_relative_idx, min=0, max=2 * c.max_relative_idx)
 
     if c.use_chain_relative:
 
@@ -586,7 +586,7 @@ def _relative_encoding(self, batch):
       max_rel_chain = c.max_relative_chain
 
       clipped_rel_chain = jnp.clip(
-          rel_sym_id + max_rel_chain, a_min=0, a_max=2 * max_rel_chain)
+          rel_sym_id + max_rel_chain, min=0, max=2 * max_rel_chain)
 
       final_rel_chain = jnp.where(entity_id_same, clipped_rel_chain,
                                   (2 * max_rel_chain + 1) *