Proposal: Support mypy --strict

[wangkuiyi]

I want to use mypy --strict to check my code that depends on MLX. However, the stub generated by pybind11-stubgen at https://github.com/ml-explore/mlx/blob/3b661b73943c8b91d45bdf55da670e8cb308c2a8/setup.py#L137C26-L137C42 does not include accurate type annotation.

In the file ~/miniforge3/envs/mlx-medusa-plus/lib/python3.10/site-packages/mlx/core/__init__.pyi, I have the following signature of mx.repeat:

def repeat(*args, **kwargs):

However, mypy --strict would expect something like the following:

repeat(array: array, repeats: int, axis: Optional[int] = None, *, stream: Union[None, Stream, Device] = None) -> array

I noticed that the help function could reveal accurate type annotations. The run of help(mx.repeat) gave me the following:

repeat(...) method of builtins.PyCapsule instance
    repeat(array: array, repeats: int, axis: Optional[int] = None, *, stream: Union[None, Stream, Device] = None) -> array

Where does help get the information from? Could pybind11-stubgen make use of such information to output accurately annotated functions? Thanks!

[wangkuiyi]

Hmm. Should I post this proposal to https://github.com/sizmailov/pybind11-stubgen instead?

[awni]

Since we switched to nanobind, the stubs in site-packages/mlx/core/__init__.pyi should have the right type information.

There is also an open issue (#1240) about getting complete typing info for MLX python packages, there are a few annotations missing here and there. I will leave that issue open and close this one as a dup / partially resolved.

[jbcoe]

I'm not seeing site-packages/mlx/core/__init__.pyi in my pip or conda installs of mlx. What do I need to do to get type stubs?

I added https://github.com/ml-explore/mlx/issues/1345 to track this.