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

nn.jit: automatic fingerprint definition for dataclass attributes #3737

Merged
merged 1 commit into from
Mar 5, 2024

Conversation

cgarciae
Copy link
Collaborator

@cgarciae cgarciae commented Mar 4, 2024

What does this PR do?

Defines a fingerprint for dataclass attributes to avoid the use of unsafe_hash.

@cgarciae cgarciae force-pushed the jit-fingerprint branch 3 times, most recently from 1fd55a1 to 2f2b755 Compare March 5, 2024 00:10
@codecov-commenter
Copy link

codecov-commenter commented Mar 5, 2024

Codecov Report

Attention: Patch coverage is 5.55556% with 51 lines in your changes are missing coverage. Please review.

Project coverage is 58.67%. Comparing base (f195a8b) to head (3ea30d7).
Report is 3 commits behind head on main.

Files Patch % Lines
flax/linen/transforms.py 5.55% 51 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main    #3737      +/-   ##
==========================================
+ Coverage   58.63%   58.67%   +0.03%     
==========================================
  Files         103      103              
  Lines       12379    12386       +7     
==========================================
+ Hits         7259     7267       +8     
+ Misses       5120     5119       -1     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

@cgarciae cgarciae force-pushed the jit-fingerprint branch 5 times, most recently from 5abdc08 to 001df6c Compare March 5, 2024 01:19
f"type '{x}' is not hashable. Unhashable fields: {unhashable_paths}"
)
try:
hash(x)
Copy link
Collaborator

Choose a reason for hiding this comment

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

would isinstance(x, typing.Hashable) work instead of a try-catch block?

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

Nope. Having a __hash__ method is not enough.

) -> tuple[type[Any], Any]:
def _fingerprint_recursive(
obj: Any, path: tuple[str, ...], seen_modules: dict[FlaxId, int]
) -> Any:
"""Creates a hashable representation for a Module by travering its structure recursively."""
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
"""Creates a hashable representation for a Module by travering its structure recursively."""
"""Creates a hashable representation for a Module by traversing its structure recursively."""

@copybara-service copybara-service bot merged commit 6760e32 into main Mar 5, 2024
21 checks passed
@copybara-service copybara-service bot deleted the jit-fingerprint branch March 5, 2024 22:25
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