[BUG] 0.32.0 wheels ship py.typed but no .pyi stubs — breaks type checking of mlx.core downstream

[dgrauet]

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:

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():

# 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.