Skip to content

typing: add type hints to mlx.core.array, linalg, and random #2565

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
awni merged 2 commits into from
Sep 4, 2025
Merged

typing: add type hints to mlx.core.array, linalg, and random #2565

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
df3725c
Add type annotations to mlx methods
XXXXRT666 Sep 4, 2025
61e9ae2
Missing list_or_scalar
XXXXRT666 Sep 4, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion python/mlx/_stub_patterns.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ mlx.core.__prefix__:
mlx.core.__suffix__:
from typing import Union
scalar: TypeAlias = Union[int, float, bool]
list_or_scalar: TypeAlias = Union[scalar, list["list_or_scalar"]]
bool_: Dtype = ...

mlx.core.distributed.__prefix__:
Expand All @@ -29,5 +30,5 @@ mlx.core.metal.__prefix__:
from typing import Sequence, Optional, Union

mlx.core.random.__prefix__:
from mlx.core import array, Dtype, Device, Stream, scalar
from mlx.core import array, Dtype, Device, Stream, scalar, float32, int32
from typing import Sequence, Optional, Union
3 changes: 3 additions & 0 deletions python/src/array.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -320,6 +320,7 @@ void init_array(nb::module_& m) {
.def_prop_ro(
"shape",
[](const mx::array& a) { return nb::cast(a.shape()); },
nb::sig("def shape(self) -> tuple[int, ...]"),
R"pbdoc(
The shape of the array as a Python tuple.

Expand Down Expand Up @@ -347,6 +348,7 @@ void init_array(nb::module_& m) {
.def(
"item",
&to_scalar,
nb::sig("def item(self) -> scalar"),
R"pbdoc(
Access the value of a scalar array.

Expand All @@ -356,6 +358,7 @@ void init_array(nb::module_& m) {
.def(
"tolist",
&tolist,
nb::sig("def tolist(self) -> list_or_scalar"),
R"pbdoc(
Convert the array to a Python :class:`list`.

Expand Down
4 changes: 4 additions & 0 deletions python/src/linalg.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -447,6 +447,8 @@ void init_linalg(nb::module_& parent_module) {
"a"_a,
nb::kw_only(),
"stream"_a = nb::none(),
nb::sig(
"def eig(a: array, *, stream: Union[None, Stream, Device] = None) -> Tuple[array, array]"),
R"pbdoc(
Compute the eigenvalues and eigenvectors of a square matrix.

Expand Down Expand Up @@ -523,6 +525,8 @@ void init_linalg(nb::module_& parent_module) {
"UPLO"_a = "L",
nb::kw_only(),
"stream"_a = nb::none(),
nb::sig(
"def eigh(a: array, UPLO: str = 'L', *, stream: Union[None, Stream, Device] = None) -> Tuple[array, array]"),
R"pbdoc(
Compute the eigenvalues and eigenvectors of a complex Hermitian or
real symmetric matrix.
Expand Down
2 changes: 1 addition & 1 deletion python/src/random.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -171,7 +171,7 @@ void init_random(nb::module_& parent_module) {
"key"_a = nb::none(),
"stream"_a = nb::none(),
nb::sig(
"def normal(shape: Sequence[int] = [], dtype: Optional[Dtype] = float32, loc: Optional[scalar, array] = None, scale: Optional[scalar, array] = None, key: Optional[array] = None, stream: Union[None, Stream, Device] = None) -> array"),
"def normal(shape: Sequence[int] = [], dtype: Optional[Dtype] = float32, loc: Union[scalar, array, None] = None, scale: Union[scalar, array, None] = None, key: Optional[array] = None, stream: Union[None, Stream, Device] = None) -> array"),
R"pbdoc(
Generate normally distributed random numbers.

Expand Down