Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Update deprecated Jax calls #35919

Merged
merged 4 commits into from
Mar 20, 2025
Merged

Update deprecated Jax calls #35919

merged 4 commits into from
Mar 20, 2025
+13 −13

Conversation

rasmi
Copy link
Contributor

@rasmi rasmi commented Jan 27, 2025

jax v0.4.27 introduced the deprecation 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 to use these new arguments.

This PR was originally written by @jakevdp.

cc @sanchit-gandhi your review would be appreciated!

@rasmi
Copy link
Contributor Author

rasmi commented Jan 27, 2025 •
edited
Loading

@sanchit-gandhi -- the test environment seems to be using the older version of Jax, despite setup.py and dependency_versions_table.py being updated. Is there a proper way to force the docker image to install the fresh set of requirements as it spins up the test environment?

@rasmi
Copy link
Contributor Author

rasmi commented Feb 6, 2025

@ArthurZucker @sanchit-gandhi -- do you know who can help with this? There are various other upcoming jax deprecations that I'd like to fix, so being able to test them correctly would be helpful.

@rasmi rasmi changed the title Update deprecated arguments in jax.numpy.clip. Update deprecated Jax calls Feb 6, 2025
@rasmi
Copy link
Contributor Author

rasmi commented Feb 7, 2025

@Rocketknight1 @amyeroberts -- please let me know how to proceed! I would love to get rid of some of these deprecated methods so that the library can be kept up-to-date.

ArthurZucker
ArthurZucker reviewed Feb 10, 2025
View reviewed changes
Copy link
Collaborator

@ArthurZucker ArthurZucker left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sounds good! The failing tests are probably due to the image docker being used, could you confirm that these pass locally? 🤗

rasmi and others added 4 commits March 17, 2025 21:44
@rasmi
Remove deprecated arguments for jax.numpy.clip.
004c2e3
@rasmi
Remove deprecated arguments for jax.numpy.clip.

Verified

This commit was created on GitHub.com and signed with GitHub’s verified signature.
GPG key ID: B5690EEEBB952194
Verified
Learn about vigilant mode
ee0aab5
@rasmi
Update jax version to 0.4.27 to 0.4.38.
4f0c56c
@rasmi @jakevdp
Avoid use of deprecated xla_bridge.get_backend().platform
Loading
Loading status checks…
d0b139d 
Co-authored-by: Jake Vanderplas <jakevdp@google.com>
@rasmi
Copy link
Contributor Author

rasmi commented Mar 18, 2025

Hi @ArthurZucker, I've rebased to main and all tests are passing! Let me know if any additional action is needed.

@rasmi rasmi requested a review from ArthurZucker March 18, 2025 13:51
@rasmi
Copy link
Contributor Author

rasmi commented Mar 18, 2025

cc @Rocketknight1 as well if you're able to review!

ArthurZucker
ArthurZucker approved these changes Mar 20, 2025
View reviewed changes
Copy link
Collaborator

@ArthurZucker ArthurZucker left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Merging!

@ArthurZucker ArthurZucker merged commit f0d5b2f into huggingface:main Mar 20, 2025
21 checks passed
@gante gante mentioned this pull request Mar 20, 2025
Merged
ArthurZucker added a commit that referenced this pull request Mar 21, 2025
@ArthurZucker
Revert "Update deprecated Jax calls (#35919)"
Loading
Loading status checks…
69b1c3b 
This reverts commit f0d5b2f.
ArthurZucker added a commit that referenced this pull request Mar 21, 2025
@ArthurZucker
Revert "Update deprecated Jax calls (#35919)"
Loading
Loading status checks…
0572513 
This reverts commit f0d5b2f.
ArthurZucker added a commit that referenced this pull request Mar 21, 2025
@ArthurZucker
Revert "Update deprecated Jax calls (#35919)" (#36880)
Loading
Loading status checks…
f19d018 
* Revert "Update deprecated Jax calls (#35919)"

This reverts commit f0d5b2f.

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

This reverts commit f0d5b2f.

* udpate
@rasmi
Copy link
Contributor Author

rasmi commented Mar 21, 2025

Hi @ArthurZucker, sorry this broke things! Do you know what went wrong? Is there a way to test this properly going forward? The CI does not seem to use an updated/custom Docker image if changes are made, so I'm not sure what the best approach is.

@ArthurZucker
Copy link
Collaborator

Yeah sorry, it was just from transformers import * or something like that

@ArthurZucker
Copy link
Collaborator

We'll try to come back to this

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@ArthurZucker ArthurZucker

Assignees
No one assigned
Labels
Flax
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.

None yet
3 participants
@rasmi @ArthurZucker @Rocketknight1