-
Notifications
You must be signed in to change notification settings - Fork 648
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
Conversation
1fd55a1
to
2f2b755
Compare
Codecov ReportAttention: Patch coverage is
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. |
5abdc08
to
001df6c
Compare
f"type '{x}' is not hashable. Unhashable fields: {unhashable_paths}" | ||
) | ||
try: | ||
hash(x) |
There was a problem hiding this comment.
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?
There was a problem hiding this comment.
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.
flax/linen/transforms.py
Outdated
) -> 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.""" |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
"""Creates a hashable representation for a Module by travering its structure recursively.""" | |
"""Creates a hashable representation for a Module by traversing its structure recursively.""" |
001df6c
to
3ea30d7
Compare
What does this PR do?
Defines a fingerprint for dataclass attributes to avoid the use of
unsafe_hash
.