From 04edad8c5908efaecdcc750d55339c02e32fcdfe Mon Sep 17 00:00:00 2001 From: Chunyang Wen Date: Fri, 7 Mar 2025 06:20:09 +0800 Subject: [PATCH] Add doc string for path (#1937) --- python/mlx/utils.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/python/mlx/utils.py b/python/mlx/utils.py index 2a3c1e660..b7173deb7 100644 --- a/python/mlx/utils.py +++ b/python/mlx/utils.py @@ -62,7 +62,7 @@ def tree_map_with_path( tree: Any, *rest: Any, is_leaf: Optional[Callable] = None, - path: Any = None, + path: Optional[Any] = None, ) -> Any: """Applies ``fn`` to the path and leaves of the Python tree ``tree`` and returns a new collection with the results. @@ -74,8 +74,9 @@ def tree_map_with_path( fn (callable): The function that processes the leaves of the tree. tree (Any): The main Python tree that will be iterated upon. rest (tuple[Any]): Extra trees to be iterated together with ``tree``. - is_leaf (callable, optional): An optional callable that returns ``True`` + is_leaf (Optional[Callable]): An optional callable that returns ``True`` if the passed object is considered a leaf or ``False`` otherwise. + path (Optional[Any]): Prefix will be added to the result. Returns: A Python tree with the new values returned by ``fn``.