Repository: huggingface/transformers
Pull request: Update deprecated Jax calls

 jax v0.4.27 introduced the [deprecation](https://jax.readthedocs.io/en/latest/changelog.html#jax-0-4-27-may-7-2024:~:text=jax.numpy.clip()%20has%20a%20new%20argument%20signature%3A%20a%2C%20a_min%2C%20and%20a_max%20are%20deprecated%20in%20favor%20of%20x%20(positional%20only)%2C%20min%2C%20and%20max%20(%2320550).) of `jax.numpy.clip`'s `a_min` and `a_max` arguments (changing to just `min` and `max`, respectively). This PR updates [all uses of jnp.clip](https://github.com/search?q=repo%3Ahuggingface%2Ftransformers+%22jnp.clip%22+%28%22a_max%22+OR+%22a_min%22%29&type=code) to use these new arguments. 

This PR was originally written by @jakevdp.

cc @sanchit-gandhi your review would be appreciated!



Changed files:
filename: setup.py
patch: @@ -121,8 +121,8 @@
     "importlib_metadata",
     "ipadic>=1.0.0,<2.0",
     "isort>=5.5.4",
-    "jax>=0.4.1,<=0.4.13",
-    "jaxlib>=0.4.1,<=0.4.13",
+    "jax>=0.4.27,<=0.4.38",
+    "jaxlib>=0.4.27,<=0.4.38",
     "jieba",
     "jinja2>=3.1.0",
     "kenlm",

filename: src/transformers/commands/env.py
patch: @@ -129,7 +129,7 @@ def run(self):
             flax_version = flax.__version__
             jax_version = jax.__version__
             jaxlib_version = jaxlib.__version__
-            jax_backend = jax.lib.xla_bridge.get_backend().platform
+            jax_backend = jax.default_backend()
 
         info = {
             "`transformers` version": version,

filename: src/transformers/dependency_versions_table.py
patch: @@ -28,8 +28,8 @@
     "importlib_metadata": "importlib_metadata",
     "ipadic": "ipadic>=1.0.0,<2.0",
     "isort": "isort>=5.5.4",
-    "jax": "jax>=0.4.1,<=0.4.13",
-    "jaxlib": "jaxlib>=0.4.1,<=0.4.13",
+    "jax": "jax>=0.4.27,<=0.4.38",
+    "jaxlib": "jaxlib>=0.4.27,<=0.4.38",
     "jieba": "jieba",
     "jinja2": "jinja2>=3.1.0",
     "kenlm": "kenlm",

filename: src/transformers/models/longt5/modeling_flax_longt5.py
patch: @@ -387,7 +387,7 @@ def _relative_position_bucket(relative_position, bidirectional=True, num_buckets
             relative_buckets += (relative_position > 0) * num_buckets
             relative_position = jnp.abs(relative_position)
         else:
-            relative_position = -jnp.clip(relative_position, a_max=0)
+            relative_position = -jnp.clip(relative_position, max=0)
         # now relative_position is in the range [0, inf)
 
         # half of the buckets are for exact increments in positions
@@ -398,7 +398,7 @@ def _relative_position_bucket(relative_position, bidirectional=True, num_buckets
         relative_position_if_large = max_exact + (
             jnp.log(relative_position / max_exact) / jnp.log(max_distance / max_exact) * (num_buckets - max_exact)
         )
-        relative_position_if_large = jnp.clip(relative_position_if_large, a_max=num_buckets - 1)
+        relative_position_if_large = jnp.clip(relative_position_if_large, max=num_buckets - 1)
 
         relative_buckets += jnp.where(is_small, relative_position, relative_position_if_large)
 
@@ -672,7 +672,7 @@ def _relative_position_bucket(relative_position, bidirectional=True, num_buckets
             relative_buckets += (relative_position > 0) * num_buckets
             relative_position = jnp.abs(relative_position)
         else:
-            relative_position = -jnp.clip(relative_position, a_max=0)
+            relative_position = -jnp.clip(relative_position, max=0)
         # now relative_position is in the range [0, inf)
 
         # half of the buckets are for exact increments in positions
@@ -683,7 +683,7 @@ def _relative_position_bucket(relative_position, bidirectional=True, num_buckets
         relative_position_if_large = max_exact + (
             jnp.log(relative_position / max_exact) / jnp.log(max_distance / max_exact) * (num_buckets - max_exact)
         )
-        relative_position_if_large = jnp.clip(relative_position_if_large, a_max=num_buckets - 1)
+        relative_position_if_large = jnp.clip(relative_position_if_large, max=num_buckets - 1)
 
         relative_buckets += jnp.where(is_small, relative_position, relative_position_if_large)
 
@@ -895,7 +895,7 @@ def _relative_position_bucket(relative_position, bidirectional=True, num_buckets
             relative_buckets += (relative_position > 0) * num_buckets
             relative_position = jnp.abs(relative_position)
         else:
-            relative_position = -jnp.clip(relative_position, a_max=0)
+            relative_position = -jnp.clip(relative_position, max=0)
         # now relative_position is in the range [0, inf)
 
         # half of the buckets are for exact increments in positions
@@ -906,7 +906,7 @@ def _relative_position_bucket(relative_position, bidirectional=True, num_buckets
         relative_position_if_large = max_exact + (
             jnp.log(relative_position / max_exact) / jnp.log(max_distance / max_exact) * (num_buckets - max_exact)
         )
-        relative_position_if_large = jnp.clip(relative_position_if_large, a_max=num_buckets - 1)
+        relative_position_if_large = jnp.clip(relative_position_if_large, max=num_buckets - 1)
 
         relative_buckets += jnp.where(is_small, relative_position, relative_position_if_large)
 

filename: src/transformers/models/t5/modeling_flax_t5.py
patch: @@ -247,7 +247,7 @@ def _relative_position_bucket(relative_position, bidirectional=True, num_buckets
             relative_buckets += (relative_position > 0) * num_buckets
             relative_position = jnp.abs(relative_position)
         else:
-            relative_position = -jnp.clip(relative_position, a_max=0)
+            relative_position = -jnp.clip(relative_position, max=0)
         # now relative_position is in the range [0, inf)
 
         # half of the buckets are for exact increments in positions
@@ -258,7 +258,7 @@ def _relative_position_bucket(relative_position, bidirectional=True, num_buckets
         relative_position_if_large = max_exact + (
             jnp.log(relative_position / max_exact) / jnp.log(max_distance / max_exact) * (num_buckets - max_exact)
         )
-        relative_position_if_large = jnp.clip(relative_position_if_large, a_max=num_buckets - 1)
+        relative_position_if_large = jnp.clip(relative_position_if_large, max=num_buckets - 1)
 
         relative_buckets += jnp.where(is_small, relative_position, relative_position_if_large)