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.

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

Add path_aware_map function #2371

Merged
merged 1 commit into from
Aug 30, 2022

Conversation

cgarciae
Copy link
Collaborator

@cgarciae cgarciae commented Aug 4, 2022

What does this PR do?

Adds a path_aware_map function to traverse_util, this function makes it easy to use create Flax with Optax'smulti_transform and masked combinators.

@codecov-commenter
Copy link

codecov-commenter commented Aug 5, 2022

Codecov Report

Merging #2371 (fbcde14) into main (f754541) will increase coverage by 9.95%.
The diff coverage is 100.00%.

@@            Coverage Diff             @@
##             main    #2371      +/-   ##
==========================================
+ Coverage   69.33%   79.28%   +9.95%     
==========================================
  Files          63       49      -14     
  Lines        5556     4875     -681     
==========================================
+ Hits         3852     3865      +13     
+ Misses       1704     1010     -694     
Impacted Files Coverage Δ
flax/traverse_util.py 99.01% <100.00%> (+0.03%) ⬆️
flax/__init__.py 100.00% <0.00%> (ø)
flax/core/lift.py 95.80% <0.00%> (ø)
flax/linen/transforms.py 94.04% <0.00%> (ø)
flax/linen/partitioning.py 82.32% <0.00%> (ø)
flax/optim/base.py
flax/optim/lars.py
flax/optim/momentum.py
flax/optim/adafactor.py
flax/optim/weight_norm.py
... and 10 more

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

@cgarciae cgarciae requested a review from marcvanzee August 5, 2022 15:31
@review-notebook-app
Copy link

Check out this pull request on  ReviewNB

See visual diffs & provide feedback on Jupyter Notebooks.


Powered by ReviewNB

@cgarciae cgarciae force-pushed the optax-partition-function branch from 0b44426 to ca10344 Compare August 5, 2022 17:53
flax/traverse_util.py Outdated Show resolved Hide resolved
flax/traverse_util.py Outdated Show resolved Hide resolved
@cgarciae cgarciae changed the title Adds variables_path_partitions function Adds path_variable_map function Aug 11, 2022
@cgarciae cgarciae changed the title Adds path_variable_map function Add path_variable_map function Aug 11, 2022
@cgarciae cgarciae changed the title Add path_variable_map function Add path_value_map function Aug 11, 2022
@cgarciae cgarciae force-pushed the optax-partition-function branch from fca9f7e to 44d1c0d Compare August 11, 2022 20:31
@cgarciae cgarciae changed the title Add path_value_map function Add path_aware_map function Aug 12, 2022
@cgarciae cgarciae force-pushed the optax-partition-function branch from 2d152c5 to f6e283b Compare August 12, 2022 20:18
@cgarciae cgarciae requested a review from jheek August 15, 2022 14:10
flax/traverse_util.py Outdated Show resolved Hide resolved
flax/traverse_util.py Outdated Show resolved Hide resolved
@cgarciae cgarciae force-pushed the optax-partition-function branch from b8e8980 to fbcde14 Compare August 26, 2022 17:47
@cgarciae cgarciae requested a review from jheek August 26, 2022 17:50
@cgarciae cgarciae force-pushed the optax-partition-function branch from fbcde14 to acb7662 Compare August 29, 2022 17:43
Copy link
Member

@jheek jheek left a comment

Choose a reason for hiding this comment

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

Looks great!

@copybara-service copybara-service bot merged commit b146e16 into google:main Aug 30, 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.

3 participants