Skip to content

Add serialization for Partial functions. Completes #2433 #2557

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
merged 3 commits into from
Dec 13, 2022

Conversation

zaxtax
Copy link
Collaborator

@zaxtax zaxtax commented Oct 25, 2022

This incorporates the feature suggestion of adding serialization support for Jax partial functions. This would benefit from a unit test, but otherwise uses the code from the earlier issue largely as is.

Fixes #2433

@cgarciae
Copy link
Collaborator

cgarciae commented Nov 2, 2022

Thanks @zaxtax! Can you add a test for this?

@zaxtax
Copy link
Collaborator Author

zaxtax commented Nov 2, 2022 via email

@zaxtax zaxtax force-pushed the partial_serialization branch from 60f8e61 to 390b42c Compare December 6, 2022 13:50
@codecov-commenter
Copy link

codecov-commenter commented Dec 6, 2022

Codecov Report

Merging #2557 (93a51c7) into main (09b6d56) will increase coverage by 0.00%.
The diff coverage is 100.00%.

@@           Coverage Diff           @@
##             main    #2557   +/-   ##
=======================================
  Coverage   81.15%   81.15%           
=======================================
  Files          51       51           
  Lines        5492     5493    +1     
=======================================
+ Hits         4457     4458    +1     
  Misses       1035     1035           
Impacted Files Coverage Δ
flax/serialization.py 94.47% <100.00%> (+0.03%) ⬆️

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

@zaxtax zaxtax requested a review from cgarciae December 6, 2022 14:25
Copy link
Collaborator

@cgarciae cgarciae left a comment

Choose a reason for hiding this comment

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

LGTM!

@copybara-service copybara-service bot merged commit 1d678b5 into google:main Dec 13, 2022
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

Support serialization for jax.tree_util.Partial
4 participants