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

[nnx] cleanup CallableProxy #3608

Merged
merged 1 commit into from
Jan 9, 2024
Merged

Conversation

cgarciae
Copy link
Collaborator

@cgarciae cgarciae commented Jan 8, 2024

What does this PR do?

  • CallableProxy now create a DelayedAccessor by default if no accessor is given.
  • Removes the creation of DelayedAccessors at many callsites CallableProxy is used.

@codecov-commenter
Copy link

codecov-commenter commented Jan 8, 2024

Codecov Report

Attention: 4 lines in your changes are missing coverage. Please review.

Comparison is base (7cbc3c1) 56.70% compared to head (b3c181d) 56.70%.

Files Patch % Lines
flax/experimental/nnx/nnx/module.py 60.00% 2 Missing ⚠️
flax/experimental/nnx/nnx/proxy_caller.py 85.71% 2 Missing ⚠️
Additional details and impacted files
@@           Coverage Diff           @@
##             main    #3608   +/-   ##
=======================================
  Coverage   56.70%   56.70%           
=======================================
  Files         100      100           
  Lines       12006    12005    -1     
=======================================
  Hits         6808     6808           
+ Misses       5198     5197    -1     

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

flax/experimental/nnx/nnx/graph_utils.py Outdated Show resolved Hide resolved
flax/experimental/nnx/nnx/module.py Outdated Show resolved Hide resolved
flax/experimental/nnx/nnx/proxy_caller.py Outdated Show resolved Hide resolved
flax/experimental/nnx/nnx/transforms.py Outdated Show resolved Hide resolved
Copy link
Collaborator

@chiamp chiamp left a comment

Choose a reason for hiding this comment

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

There is also a CI failure

@cgarciae
Copy link
Collaborator Author

cgarciae commented Jan 9, 2024

Thanks! Fixed.

@@ -330,10 +330,10 @@ def apply(
accessesor = DelayedAccessor()
Copy link
Collaborator

Choose a reason for hiding this comment

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

still some typos here: accessesor -> accessor

Comment on lines 406 to 409
accessesor: DelayedAccessor, *args, **kwargs
) -> tuple[tp.Any, M]:
module = self.clone()
fn = accessesor(module)
Copy link
Collaborator

Choose a reason for hiding this comment

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

and here

@copybara-service copybara-service bot merged commit f63808c into main Jan 9, 2024
21 checks passed
@copybara-service copybara-service bot deleted the nnx-improve-callable-proxy branch January 9, 2024 20:19
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