Revert "Update deprecated Jax calls (#35919)" (#36880)

* Revert "Update deprecated Jax calls (#35919)" This reverts commit f0d5b2f. * Revert "Update deprecated Jax calls (#35919)" This reverts commit f0d5b2f. * udpate

Author: ArthurZucker

src/transformers/commands/env.py

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

src/transformers/models/longt5/modeling_flax_longt5.py

@@ -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, max=0)
+ relative_position = -jnp.clip(relative_position, a_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, max=num_buckets - 1)
+ relative_position_if_large = jnp.clip(relative_position_if_large, a_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, max=0)
+ relative_position = -jnp.clip(relative_position, a_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, max=num_buckets - 1)
+ relative_position_if_large = jnp.clip(relative_position_if_large, a_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, max=0)
+ relative_position = -jnp.clip(relative_position, a_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, max=num_buckets - 1)
+ relative_position_if_large = jnp.clip(relative_position_if_large, a_max=num_buckets - 1)
 
  relative_buckets += jnp.where(is_small, relative_position, relative_position_if_large)
 

src/transformers/models/t5/modeling_flax_t5.py

@@ -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, max=0)
+ relative_position = -jnp.clip(relative_position, a_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, max=num_buckets - 1)
+ relative_position_if_large = jnp.clip(relative_position_if_large, a_max=num_buckets - 1)
 
  relative_buckets += jnp.where(is_small, relative_position, relative_position_if_large)