zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-12-16 01:49:05 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

Some fixes to typing (#1371)

Browse Source 
* some fixes to typing

* fix module reference

* comment
This commit is contained in:
Awni Hannun
2024-08-28 11:16:19 -07:00
committed by GitHub
parent bd47e1f066
commit 291cf40aca
15 changed files with 152 additions and 145 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
python/mlx/utils.py
View File

@@ -1,10 +1,10 @@
# Copyright © 2023 Apple Inc.
from collections import defaultdict
from typing import Any, Callable, Tuple
from typing import Any, Callable, Optional, Tuple
def tree_map(
fn: Callable, tree: Any, *rest: Tuple[Any], is_leaf: Callable = None
fn: Callable, tree: Any, *rest: Any, is_leaf: Optional[Callable] = None
) -> Any:
"""Applies ``fn`` to the leaves of the Python tree ``tree`` and
returns a new collection with the results.
@@ -59,8 +59,8 @@ def tree_map(
def tree_map_with_path(
fn: Callable,
tree: Any,
*rest: Tuple[Any],
is_leaf: Callable = None,
*rest: Any,
is_leaf: Optional[Callable] = None,
path: Any = None,
) -> Any:
"""Applies ``fn`` to the path and leaves of the Python tree ``tree`` and