[BUG] 0.32.0 wheels ship py.typed but no .pyi stubs — breaks type checking of mlx.core downstream
envgap__ml-explore__mlx-3916
01 / FAILURE SIGNATURE
As reported upstream
error[unresolved-import]: Cannot resolve imported module `mlx.core`
Not a benchmark task.
- In a clean container the reported failure did not reproduce, or the known fix did not make the project run.
02 / ENVIRONMENT RECIPE
- Base commit
61a4867be62594aa2a243910b401ac0bf3e02525- Manifest
setup.py- Reproduce
Awaiting issue-specific recipe- Run under trace
Awaiting a meaningful runtime command
03 / ORIGINAL ISSUE TEXT
ml-explore/mlx #3916 · read the original issue
### Summary
The macOS arm64 wheels for **0.32.0** ship `mlx/py.typed` but contain **no `.pyi` files at all**. The 0.31.2 wheels contain 8 (`mlx/core/__init__.pyi`, `cuda.pyi`, `distributed.pyi`, `fast.pyi`, …).
This is the worst of both worlds for downstream type checking: `py.typed` asserts the package is typed, so checkers trust it and do not fall back to inference or bundled third-party stubs — but there is nothing to resolve `mlx.core` against, since it is a compiled extension.
### Reproduction
No install needed — inspect the published wheels directly:
```python
import urllib.request, json, io, zipfile
for ver in ("0.31.2", "0.32.0"):
d = json.load(urllib.request.urlopen(f"https://pypi.org/pypi/mlx/{ver}/json"))
url = next(f["url"] for f in d["urls"]
if f["filename"].endswith("cp312-cp312-macosx_14_0_arm64.whl"))
names = zipfile.ZipFile(io.BytesIO(urllib.request.urlopen(url).read())).namelist()
pyi = [n for n in names if n.endswith(".pyi")]
print(ver, "pyi:", len(pyi), "py.typed:", "mlx/py.typed" in names, "files:", len(names))
```
```
0.31.2 pyi: 8 py.typed: True files: 47
0.32.0 pyi: 0 py.typed: True files: 39
```
Same result for `cp311`. (Checked on 2026-07-25; 0.32.0 is the latest release.)
### Downstream impact
Any type checker fails on every import. With `ty`:
```
error[unresolved-import]: Cannot resolve imported module `mlx.core`
--> src/…/attention.py:3
|
3 | import mlx.core as mx
```
Runtime is completely unaffected — this is packaging-only. In our project we had to constrain the **dev** environment to `mlx<0.32` to keep the type-check job green, which in turn blocks adopting any 0.32-only API (e.g. the `mx.fast.metal_kernel` math-mode option from #3728), since a `mlx>=0.32` runtime floor would make the dev extra uninstallable.
### Possible cause
Stub generation does not appear to have been removed deliberately — `setup.py` still runs it at the end of `build_ext.run()`:
```python
# Build type stubs.
build_temp = Path(self.build_temp) / ext.name
subprocess.run(
["cmake", "--install", build_temp, "--component", "core_stub"],
check=True,
)
```
So the stubs are most likely generated but then dropped before the wheel is assembled. One hypothesis worth checking first: the two-stage release build (`MLX_BUILD_STAGE=1` for `mlx`, stage 2 for `mlx-metal`) combined with the CI rework in #3491, which moved `python setup.py clean --all` to run *after* `python -m build -w` and added `uv venv --clear` before the build. Something in that ordering could remove or hide the installed stubs relative to the directory that gets packaged. I have not bisected this — treat it as a starting point, not a diagnosis.
Prior art for the same class of regression: #2398 ("Fix including stubs in wheel") and #2942 ("Fix stubgen").
### Expected
0.32.x wheels ship `mlx/core/*.pyi` again, as 0.31.2 did. Failing that, dropping `mlx/py.typed` would at least let downstream checkers fall back gracefully instead of hard-failing — but restoring the stubs is clearly the better outcome.
A regression test in the release job — assert the built wheel contains `mlx/core/__init__.pyi` before upload — would prevent a third recurrence.
04 / LABELS
Labels from the report text only; not yet run
No supported category has been assigned.
Label rules and the text that matched
[]