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] add PathContains Filter #4011

Merged
merged 1 commit into from
Jun 21, 2024
Merged

[nnx] add PathContains Filter #4011

merged 1 commit into from
Jun 21, 2024

Conversation

cgarciae
Copy link
Collaborator

What does this PR do?

Adds the PathContains Filter which lets users select paths that contain a specific key. Example

class Model(nnx.Module):
  def __init__(self, rngs):
    self.backbone = nnx.Linear(2, 3, rngs=rngs)
    self.head = nnx.Linear(3, 10, rngs=rngs)

model = Model(nnx.Rngs(0))

head_state = nnx.state(model, nnx.PathContains('head'))

assert 'head' in head_state
assert 'backbone' not in head_state

@copybara-service copybara-service bot merged commit 4d9081e into main Jun 21, 2024
21 checks passed
@copybara-service copybara-service bot deleted the nnx-path-filter branch June 21, 2024 12:13
cgarciae pushed a commit that referenced this pull request Jun 24, 2024
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.

2 participants