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

make flax.core.copy add_or_replace optional #3241

Merged
merged 4 commits into from
Aug 4, 2023

Conversation

PhilipVinc
Copy link
Contributor

It makes the API easier to use when you just want to have a copy of a dictionary.

ignore .envrc (direnv files)


format


add test


fix
flax/core/frozen_dict.py Outdated Show resolved Hide resolved
flax/core/frozen_dict.py Outdated Show resolved Hide resolved
Comment on lines 251 to 252
if add_or_replace is not None:
new_dict.update(add_or_replace)
Copy link
Collaborator

Choose a reason for hiding this comment

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

Suggested change
if add_or_replace is not None:
new_dict.update(add_or_replace)
if isinstance(add_or_replace, dict):
add_or_replace = jax.tree_map(lambda x: x, add_or_replace)
new_dict.update(add_or_replace)

Copy link
Contributor Author

Choose a reason for hiding this comment

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

I'm not sure I understand the logic here.
Using your MappingProxy idea, I can just remove this change, and wouldn't need to touch those lines.

@cgarciae
Copy link
Collaborator

cgarciae commented Aug 1, 2023

Thanks @PhilipVinc! Added some suggestions.

cc @chiamp

PhilipVinc and others added 2 commits August 1, 2023 10:15
Co-authored-by: Cristian Garcia <cgarcia.e88@gmail.com>
pyink
@PhilipVinc
Copy link
Contributor Author

thanks @cgarciae . I implemented your suggestions except for the one above where I just removed by edits. Let me know if this looks good for you.

@chiamp
Copy link
Collaborator

chiamp commented Aug 2, 2023

This looks great, thanks @PhilipVinc!

@copybara-service copybara-service bot merged commit d883d34 into google:main Aug 4, 2023
16 checks passed
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