Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

Sign up
Appearance settings

Deprecate Flax support #12151

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.

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

Merged
sayakpaul merged 15 commits into main from deprecate-jax
Aug 26, 2025
Merged

Deprecate Flax support #12151

sayakpaul merged 15 commits into main from deprecate-jax
Aug 26, 2025

Conversation

Copy link
Member

@sayakpaul sayakpaul commented Aug 15, 2025

651961 reacted with rocket emoji Davidnet reacted with eyes emoji

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

@sayakpaul sayakpaul marked this pull request as ready for review August 15, 2025 03:55
@sayakpaul sayakpaul changed the title (削除) [wip] Deprecate Flax support (削除ここまで) (追記) Deprecate Flax support (追記ここまで) Aug 15, 2025
@@ -151,6 +156,11 @@ class FlaxAttention(nn.Module):
dtype: jnp.dtype = jnp.float32

def setup(self):
logger.warning(
Copy link
Member Author

@sayakpaul sayakpaul Aug 15, 2025

Choose a reason for hiding this comment

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

Thought of adding it to all the public Flax classes, can remove.

Copy link
Collaborator

@DN6 DN6 left a comment

Choose a reason for hiding this comment

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

LGTM 👍🏽

Copy link
Member Author

@DN6 should we consider removing JAX and Flax from setup.py as well?

DN6 reacted with thumbs up emoji

Copy link
Member Author

Over to you, @yiyixuxu :)

Copy link
Collaborator

@yiyixuxu yiyixuxu left a comment

Choose a reason for hiding this comment

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

thank yoU!

Copy link
Member Author

Failing test is unrelated.

@sayakpaul sayakpaul merged commit 532f41c into main Aug 26, 2025
18 of 19 checks passed
@sayakpaul sayakpaul deleted the deprecate-jax branch August 26, 2025 07:58
Copy link
Collaborator

DN6 commented Aug 26, 2025

@sayakpaul Can do in a follow up, but we also need to make torch a required dep.

sayakpaul reacted with thumbs up emoji

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

@DN6 DN6 DN6 approved these changes

@yiyixuxu yiyixuxu yiyixuxu approved these changes

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

Successfully merging this pull request may close these issues.

AltStyle によって変換されたページ (->オリジナル) /