[BUG] Stubs generation on initial setup

[levje]

Describe the bug
Hi, I am simply trying to get a simple setup locally of MLX to hopefully contribute to the repo in the upcoming months. I was following the docs listing the install procedure for the Python API, and I am not able to generate the stubs.

Am I doing something wrong or does anyone have an idea of where I should look into to get my setup properly working?

python setup.py generate_stubs
Processing /Users/user/Documents/projects/mlx
  Installing build dependencies ... done
  Getting requirements to build wheel ... done
  Installing backend dependencies ... done
  Preparing metadata (pyproject.toml) ... done
Requirement already satisfied: pre-commit in /opt/homebrew/Caskroom/miniconda/base/envs/mlx/lib/python3.10/site-packages (3.7.0)
Requirement already satisfied: cfgv>=2.0.0 in /opt/homebrew/Caskroom/miniconda/base/envs/mlx/lib/python3.10/site-packages (from pre-commit) (3.4.0)
Requirement already satisfied: identify>=1.0.0 in /opt/homebrew/Caskroom/miniconda/base/envs/mlx/lib/python3.10/site-packages (from pre-commit) (2.5.36)
Requirement already satisfied: nodeenv>=0.11.1 in /opt/homebrew/Caskroom/miniconda/base/envs/mlx/lib/python3.10/site-packages (from pre-commit) (1.8.0)
Requirement already satisfied: pyyaml>=5.1 in /opt/homebrew/Caskroom/miniconda/base/envs/mlx/lib/python3.10/site-packages (from pre-commit) (6.0.1)
Requirement already satisfied: virtualenv>=20.10.0 in /opt/homebrew/Caskroom/miniconda/base/envs/mlx/lib/python3.10/site-packages (from pre-commit) (20.25.3)
Requirement already satisfied: setuptools in /opt/homebrew/Caskroom/miniconda/base/envs/mlx/lib/python3.10/site-packages (from nodeenv>=0.11.1->pre-commit) (68.2.2)
Requirement already satisfied: distlib<1,>=0.3.7 in /opt/homebrew/Caskroom/miniconda/base/envs/mlx/lib/python3.10/site-packages (from virtualenv>=20.10.0->pre-commit) (0.3.8)
Requirement already satisfied: filelock<4,>=3.12.2 in /opt/homebrew/Caskroom/miniconda/base/envs/mlx/lib/python3.10/site-packages (from virtualenv>=20.10.0->pre-commit) (3.13.4)
Requirement already satisfied: platformdirs<5,>=3.9.1 in /opt/homebrew/Caskroom/miniconda/base/envs/mlx/lib/python3.10/site-packages (from virtualenv>=20.10.0->pre-commit) (4.2.0)
Building wheels for collected packages: mlx
  Building wheel for mlx (pyproject.toml) ... done
  Created wheel for mlx: filename=mlx-0.11.1.dev20240420+b0012cd-cp310-cp310-macosx_14_0_arm64.whl size=32740014 sha256=ab0491190dac5b05ee40d419cbeae4050359a94f80c5b2f460ec0c6c23ca9772
  Stored in directory: /private/var/folders/_z/0knfg7qd17b4xt9bvl1dyb8m0000gn/T/pip-ephem-wheel-cache-i9ga46ch/wheels/c4/30/84/9609db98b56ff07bfe98541b526d7164bbb330a43ba658252e
Successfully built mlx
Installing collected packages: mlx
  Attempting uninstall: mlx
    Found existing installation: mlx 0.11.1.dev20240420+b0012cd
    Uninstalling mlx-0.11.1.dev20240420+b0012cd:
      Successfully uninstalled mlx-0.11.1.dev20240420+b0012cd
Successfully installed mlx-0.11.1.dev20240420+b0012cd
running generate_stubs
Module "mlx.core" ..
  - importing ..
  - analyzing ..
Traceback (most recent call last):
  File "/opt/homebrew/Caskroom/miniconda/base/envs/mlx/lib/python3.10/runpy.py", line 196, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/opt/homebrew/Caskroom/miniconda/base/envs/mlx/lib/python3.10/runpy.py", line 86, in _run_code
    exec(code, run_globals)
  File "/opt/homebrew/Caskroom/miniconda/base/envs/mlx/lib/python3.10/site-packages/nanobind/stubgen.py", line 1391, in <module>
    main()
  File "/opt/homebrew/Caskroom/miniconda/base/envs/mlx/lib/python3.10/site-packages/nanobind/stubgen.py", line 1362, in main
    sg.put(mod_imported)
  File "/opt/homebrew/Caskroom/miniconda/base/envs/mlx/lib/python3.10/site-packages/nanobind/stubgen.py", line 811, in put
    self.put(child, name=name, parent=value)
  File "/opt/homebrew/Caskroom/miniconda/base/envs/mlx/lib/python3.10/site-packages/nanobind/stubgen.py", line 818, in put
    self.put_type(value, name)
  File "/opt/homebrew/Caskroom/miniconda/base/envs/mlx/lib/python3.10/site-packages/nanobind/stubgen.py", line 515, in put_type
    self.put_docstr(docstr)
  File "/opt/homebrew/Caskroom/miniconda/base/envs/mlx/lib/python3.10/site-packages/nanobind/stubgen.py", line 281, in put_docstr
    docstr = textwrap.dedent(docstr).strip()
  File "/opt/homebrew/Caskroom/miniconda/base/envs/mlx/lib/python3.10/textwrap.py", line 438, in dedent
    text = _whitespace_only_re.sub('', text)
TypeError: expected string or bytes-like object
rm: python/mlx/core/mlx.pyi: No such file or directory
Module "mlx.core" ..
  - importing ..
  - analyzing ..
Traceback (most recent call last):
  File "/opt/homebrew/Caskroom/miniconda/base/envs/mlx/lib/python3.10/runpy.py", line 196, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/opt/homebrew/Caskroom/miniconda/base/envs/mlx/lib/python3.10/runpy.py", line 86, in _run_code
    exec(code, run_globals)
  File "/opt/homebrew/Caskroom/miniconda/base/envs/mlx/lib/python3.10/site-packages/nanobind/stubgen.py", line 1391, in <module>
    main()
  File "/opt/homebrew/Caskroom/miniconda/base/envs/mlx/lib/python3.10/site-packages/nanobind/stubgen.py", line 1362, in main
    sg.put(mod_imported)
  File "/opt/homebrew/Caskroom/miniconda/base/envs/mlx/lib/python3.10/site-packages/nanobind/stubgen.py", line 811, in put
    self.put(child, name=name, parent=value)
  File "/opt/homebrew/Caskroom/miniconda/base/envs/mlx/lib/python3.10/site-packages/nanobind/stubgen.py", line 818, in put
    self.put_type(value, name)
  File "/opt/homebrew/Caskroom/miniconda/base/envs/mlx/lib/python3.10/site-packages/nanobind/stubgen.py", line 515, in put_type
    self.put_docstr(docstr)
  File "/opt/homebrew/Caskroom/miniconda/base/envs/mlx/lib/python3.10/site-packages/nanobind/stubgen.py", line 281, in put_docstr
    docstr = textwrap.dedent(docstr).strip()
  File "/opt/homebrew/Caskroom/miniconda/base/envs/mlx/lib/python3.10/textwrap.py", line 438, in dedent
    text = _whitespace_only_re.sub('', text)
TypeError: expected string or bytes-like object

To Reproduce

# The previous steps I've done (which all worked fine)
conda create -n mlx-env python=3.10
pip install git+https://github.com/wjakob/nanobind.git
env CMAKE_BUILD_PARALLEL_LEVEL="" pip install -e .
pip install ".[testing]"
python -m unittest discover python/tests
pip install ".[dev]"

# The problematic line
python setup.py generate_stubs 

Expected behavior
I would just expect the stubs to be installing properly since I'm on a fresh setup following the documentation. Or either the documentation is to be adapted to show additional steps that I'd need to take in such a scenario.

Desktop (please complete the following information):

• OS Version: MacOS 14.3.1
• Python 3.10.14
• conda 23.7.3

Thank you,
Have a great weekend!

[awni]

I just tried generating the stubs, it works fine for me..

I am not sure why it's failing to generate the stubs for you. Maybe the latest nanobind + MLX is broken there.

You could try installing nanobind:

pip install git+https://github.com/wjakob/nanobind.git@2f04eac452a6d9142dedb957701bdb20125561e4

And then rebuild MLX and generate the stubs

[levje]

Thank you for your fast response.

I retried yesterday with Python 3.10, with the nanobind hash you provided, a new environment and a clean repo. It still gave me the same result.

I gave it a try with Python 3.8 just out of curiosity and it surprisingly worked (after installing manually pip install typing_extensions for some reason). I tried again with Python 3.11 to make sure, and it also worked.

Not too sure if it's the combination of the Python version with nanobind or whether it's a setup problem on my side, but I have it working now. So, I'll close the issue.