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

Adding missing type hints (#1243)

Browse Source 
* added type hints for `run`, `tree_map` and `tree_map_with_path`

* fix lint

---------

Co-authored-by: Awni Hannun <awni@apple.com>
This commit is contained in:
toji 2024-07-23 19:59:38 +05:30 committed by GitHub
parent 6307d166eb
commit 6768c6a54a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 12 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
python/mlx/extension.py
View File

@ -77,7 +77,7 @@ class CMakeBuild(build_ext):
["cmake", "--build", ".", *build_args], cwd=build_temp, check=True
)
def run(self):
def run(self) -> None:
super().run()
# Based on https://github.com/pypa/setuptools/blob/main/setuptools/command/build_ext.py#L102

13
python/mlx/utils.py
View File

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