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,
*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``.