search

ml-explore / mlx

MLX: An array framework for Apple silicon
https://ml-explore.github.io/mlx/
MIT License
15.83k stars 903 forks source link

Proposal: Support mypy --strict #746

Open wangkuiyi opened 4 months ago

wangkuiyi commented 4 months ago

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 commented 4 months ago

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