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

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.

Thanks for reporting the issue, I'm not familiar with python's type system, what would be the best command to verify the correctness of type stubs so we can run it in CI?

uv pip install mlx
uv run ????

You can run uv check to verify correctness of the type stubs.

@zcbenz

@zcbenz thanks for looking at this — happy to give you a concrete recipe.

One small clarification so CI doesn't end up green by accident: uv check verifies a project's metadata/lockfile rather than type stubs, so it will pass either way here. The checks below are aimed specifically at the .pyi files.

The cheap gate: assert the stubs are in the wheel

This is the one that would actually have caught the regression, and it needs no type checker:

python -m zipfile -l dist/mlx-*.whl | grep -E 'mlx/core/.*\.pyi'

Or as a hard CI check right after the frontend wheel is built:

python - <<'EOF'
import glob, sys, zipfile
whl = glob.glob("dist/mlx-*.whl")[0]
names = zipfile.ZipFile(whl).namelist()
missing = [f for f in ("mlx/core/__init__.pyi", "mlx/core/fast.pyi",
                       "mlx/core/random.pyi", "mlx/core/linalg.pyi")
           if f not in names]
if missing:
    sys.exit(f"{whl}: missing stubs {missing}")
print(f"{whl}: stubs OK")
EOF

For reference, the frontend wheels contain (same platform tag, cp312-macosx_15_0_arm64):

version	`mlx/py.typed`	`mlx/core/*.pyi`
0.31.2	yes	8 files (`__init__`, `cuda`, `distributed`, `fast`, `fft`, `linalg`, `metal`, `random`)
0.32.0	yes	0 files

The end-to-end check: type-check a downstream snippet

If you want CI to verify that the shipped package is actually usable by a type checker (this is what breaks for downstream projects):

uv pip install mlx mypy
cat > stub_probe.py <<'EOF'
import mlx.core as mx

a: mx.array = mx.zeros((2, 3), dtype=mx.float32)
b: mx.array = mx.fast.rms_norm(a, mx.ones((3,)), 1e-5)
reveal_type(mx.zeros)
EOF
mypy stub_probe.py

Verified A/B in a clean venv (Python 3.12, macOS arm64), only the mlx version differs:

# mlx==0.32.0
stub_probe.py:1: error: Cannot find implementation or library stub for module named "mlx.core"  [import-not-found]
stub_probe.py:5: note: Revealed type is "Any"
Found 1 error in 1 file (checked 1 source file)

# mlx==0.31.2
stub_probe.py:5: note: Revealed type is "def (shape: int | typing.Sequence[int], dtype: mlx.core.Dtype | None =, *, stream: None | mlx.core.Stream | mlx.core.Device =) -> mlx.core.array"
Success: no issues found in 1 source file

Any checker works here (pyright, ty, …); mypy just has the least setup. Note that because py.typed is shipped, the missing stubs are worse than having no typing marker at all: downstream checkers trust the package and silently resolve everything through the extension module.

(python -m mypy.stubtest mlx.core is the tool for verifying stubs match the runtime, but it already reports a handful of pre-existing nits on 0.31.2, so I wouldn't make it a blocking gate — the two checks above are enough for this regression.)

Where it likely regressed

The stub machinery itself is unchanged between 0.31.2 and 0.32.0 — python/src/CMakeLists.txt is byte-identical, and setup.py still has both nanobind_add_stub/--component core_stub and package_data = {"mlx.core": ["*.pyi"]}.

What changed is how release wheels are built. In .github/actions/build-wheel/action.yml the frontend step now runs:

python setup.py clean --all
MLX_BUILD_STAGE=1 python -m build -w

The stubs are generated in CMakeBuild.run() (i.e. during build_ext), which setuptools runs after build_py — and build_py is what collects package_data from python/mlx/core/*.pyi. So on a tree that already has stubs from a previous build the files get picked up, but on a freshly cleaned tree (clean --all + a fresh CI checkout) build_py runs before they exist and the wheel ships without them. That would explain why local builds look fine while the released wheels don't.

Haven't tested a fix, so treat that as a hypothesis — but if it's right, generating the stubs before build_py (or adding the .pyi files via a build_py/bdist_wheel hook after build_ext) plus the zip assertion above should close it for good.

Happy to open a PR if that's useful.

Thanks for looking into this, I had included a fix in https://github.com/ml-explore/mlx/pull/3818/changes#diff-60f61ab7a8d1910d86d9fda2261620314edcae5894d5aaa236b821c7256badd7R193 and I could verify the type stubs were added to the wheels https://github.com/ml-explore/mlx/actions/runs/31754270684/job/94626578013. We just have too much type errors to run uv check in CI.

I think there are two issues here:

missing type stubs in the latest build, and
type errors in the codebase.

uv check will error on both of these, but may be possible to parse its output to separate the two issues. Otherwise a simpler check for the presence of the type stubs?

build(deps): require mlx>=0.32.1 in the dev extra (#54)

Navigation Menu

Repository navigation

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

Summary

Reproduction

Downstream impact

Possible cause

Expected

The cheap gate: assert the stubs are in the wheel

The end-to-end check: type-check a downstream snippet

Where it likely regressed

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

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

Description

Summary

Reproduction

Downstream impact

Possible cause

Expected

Activity

zcbenz commented on Aug 3, 2026

alec-bike commented on Aug 13, 2026

dgrauet commented on Aug 14, 2026

The cheap gate: assert the stubs are in the wheel

The end-to-end check: type-check a downstream snippet

Where it likely regressed

zcbenz commented on Aug 15, 2026

alec-bike commented on Aug 16, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions