← All tasks
pythonml-explore/mlx #3916Not a task: not reproduced

[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.
Continue on GitHub ↗

04 / LABELS

Labels from the report text only; not yet run

No supported category has been assigned.

Label rules and the text that matched
[]