From 921cec38c096f61d759398ce4794c236ae327f4f Mon Sep 17 00:00:00 2001 From: ayaangazali Date: Fri, 24 Jul 2026 01:01:16 -0700 Subject: [PATCH] Fix filter_and_map docstring argument order for filter_fn and is_leaf_fn --- python/mlx/nn/layers/base.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/python/mlx/nn/layers/base.py b/python/mlx/nn/layers/base.py index 9802c6c311..c2bcf7ff36 100644 --- a/python/mlx/nn/layers/base.py +++ b/python/mlx/nn/layers/base.py @@ -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