zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-06-24 17:31:16 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Add doc string for path (#1937)

Browse Source
This commit is contained in:
Chunyang Wen 2025-03-07 06:20:09 +08:00 committed by GitHub
parent 392b3060b0
commit 04edad8c59
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 3 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
python/mlx/utils.py
View File

@ -62,7 +62,7 @@ def tree_map_with_path(
tree: Any, tree: Any,
*rest: Any, *rest: Any,
is_leaf: Optional[Callable] = None, is_leaf: Optional[Callable] = None,
path: Any = None, path: Optional[Any] = None,
) -> Any: ) -> Any:
"""Applies ``fn`` to the path and leaves of the Python tree ``tree`` and """Applies ``fn`` to the path and leaves of the Python tree ``tree`` and
returns a new collection with the results. 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. fn (callable): The function that processes the leaves of the tree.
tree (Any): The main Python tree that will be iterated upon. tree (Any): The main Python tree that will be iterated upon.
rest (tuple[Any]): Extra trees to be iterated together with ``tree``. 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. if the passed object is considered a leaf or ``False`` otherwise.
path (Optional[Any]): Prefix will be added to the result.
Returns: Returns:
A Python tree with the new values returned by ``fn``. A Python tree with the new values returned by ``fn``.