Skip to content

Fix filter_and_map docstring argument order for filter_fn and is_leaf_fn #3906

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.

Sign up for GitHub

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

Open
ayaangazali wants to merge 1 commit into
base: main
Choose a base branch
from
+4 −4
Open

Fix filter_and_map docstring argument order for filter_fn and is_leaf_fn #3906

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 4 additions & 4 deletions python/mlx/nn/layers/base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -255,13 +255,13 @@ def filter_and_map(
but it can also be used to extract any subset of the module's parameters.

Args:
filter_fn (Callable): Given a value, the key in which it is found
and the containing module, decide whether to keep the value or
filter_fn (Callable): Given the containing module, the key in which
it is found and the value, decide whether to keep the value or
drop it.
map_fn (Callable, optional): Optionally transform the value before
returning it.
is_leaf_fn (Callable, optional): Given a value, the key in which it
is found and the containing module decide if it is a leaf.
is_leaf_fn (Callable, optional): Given the containing module, the
key in which it is found and the value decide if it is a leaf.

Returns:
A dictionary containing the contents of the module recursively filtered
Expand Down