From 9ce86383926e0dc22899fa44165875477ffd6ccc Mon Sep 17 00:00:00 2001 From: Chaoran Yu Date: Mon, 24 Nov 2025 14:10:44 -0800 Subject: [PATCH 1/3] Update docstring for apply_fn argument Clarified the apply_fn argument in the docstring to include details about its parameters. --- python/mlx/nn/layers/base.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python/mlx/nn/layers/base.py b/python/mlx/nn/layers/base.py index 1ae9272799..61fcaf7ac3 100644 --- a/python/mlx/nn/layers/base.py +++ b/python/mlx/nn/layers/base.py @@ -407,7 +407,9 @@ def apply_to_modules(self, apply_fn: Callable[[str, Module], Any]) -> Module: instance). Args: - apply_fn (Callable): The function to apply to the modules. + apply_fn (Callable): The function to apply to the modules. It takes two parameters. + The first parameter is a string representing the path of the module in the hierarchy (e.g. "transformer.layers.0"). + The second parameter is the actual module object Returns: The module instance after updating submodules. From 01b76d363bbae80882126b8f9cfb9d361acb4cb3 Mon Sep 17 00:00:00 2001 From: Chaoran Yu Date: Mon, 24 Nov 2025 14:15:23 -0800 Subject: [PATCH 2/3] Reformatted code --- python/mlx/nn/layers/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/mlx/nn/layers/base.py b/python/mlx/nn/layers/base.py index 61fcaf7ac3..32bff9d4f5 100644 --- a/python/mlx/nn/layers/base.py +++ b/python/mlx/nn/layers/base.py @@ -407,7 +407,7 @@ def apply_to_modules(self, apply_fn: Callable[[str, Module], Any]) -> Module: instance). Args: - apply_fn (Callable): The function to apply to the modules. It takes two parameters. + apply_fn (Callable): The function to apply to the modules. It takes two parameters. The first parameter is a string representing the path of the module in the hierarchy (e.g. "transformer.layers.0"). The second parameter is the actual module object From 80c3be97c8355ddd79e8f6357ac80b5bd191381e Mon Sep 17 00:00:00 2001 From: Awni Hannun Date: Tue, 25 Nov 2025 12:30:42 -0800 Subject: [PATCH 3/3] nits --- python/mlx/nn/layers/base.py | 7 ++++--- python/src/transforms.cpp | 28 ++++++++++++++++++++++++---- 2 files changed, 28 insertions(+), 7 deletions(-) diff --git a/python/mlx/nn/layers/base.py b/python/mlx/nn/layers/base.py index 32bff9d4f5..ce009d9f83 100644 --- a/python/mlx/nn/layers/base.py +++ b/python/mlx/nn/layers/base.py @@ -407,9 +407,10 @@ def apply_to_modules(self, apply_fn: Callable[[str, Module], Any]) -> Module: instance). Args: - apply_fn (Callable): The function to apply to the modules. It takes two parameters. - The first parameter is a string representing the path of the module in the hierarchy (e.g. "transformer.layers.0"). - The second parameter is the actual module object + apply_fn (Callable): The function to apply to the modules which + takes two parameters. The first parameter is the string path of + the module (e.g. ``"model.layers.0.linear"``). The second + parameter is the module object. Returns: The module instance after updating submodules. diff --git a/python/src/transforms.cpp b/python/src/transforms.cpp index 12aa641f89..0530fa089b 100644 --- a/python/src/transforms.cpp +++ b/python/src/transforms.cpp @@ -1238,8 +1238,18 @@ void init_transforms(nb::module_& m) { same in number, shape, and type as the inputs of ``fun`` (i.e. the ``primals``). Returns: - list(array): A list of the Jacobian-vector products which - is the same in number, shape, and type of the inputs to ``fun``. + tuple(list(array), list(array)): A tuple with the outputs of + ``fun`` in the first position and the Jacobian-vector products + in the second position. + + Example: + + .. code-block:: python + + import mlx.core as mx + + outs, jvps = mx.jvp(mx.sin, (mx.array(1.0),), (mx.array(1.0),)) + )pbdoc"); m.def( "vjp", @@ -1277,8 +1287,18 @@ void init_transforms(nb::module_& m) { same in number, shape, and type as the outputs of ``fun``. Returns: - list(array): A list of the vector-Jacobian products which - is the same in number, shape, and type of the outputs of ``fun``. + tuple(list(array), list(array)): A tuple with the outputs of + ``fun`` in the first position and the vector-Jacobian products + in the second position. + + Example: + + .. code-block:: python + + import mlx.core as mx + + outs, vjps = mx.vjp(mx.sin, (mx.array(1.0),), (mx.array(1.0),)) + )pbdoc"); m.def( "value_and_grad",