ml-explore / ml-explore/mlx

mx.compile inlines float scalar constants with 7 significant digits, so compiled results differ from eager by 1 ulp

Open Beginner friendly
#4,503 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
C++
Stars
28.5k
Forks
2.3k
Avg merge
3d 8h
Merged PRs (30d)
62

Description

Describe the bug

A compiled function that captures a Python float (or a 0-d constant array) prints it into the generated kernel source with std::setprecision(std::numeric_limits<float>::digits10 + 1) (mlx/backend/common/compiled.h, print_float_constant). That is 7 significant digits, but round-tripping a float32 needs max_digits10 (9), so constants like 1/3, 128 ** -0.5 or 1/sqrt(2) are parsed back by the Metal compiler as a neighbouring float. The compiled function then differs from eager execution by 1 ulp on almost every element. Constants whose 7-digit decimal form happens to round-trip (0.3, 0.1, 1e-6) are unaffected, which makes this easy to miss in tests.

The same off-by-two applies to double (digits10 + 1 = 16, max_digits10 = 17).

To Reproduce

import mlx.core as mx

mx.random.seed(0)
x = mx.random.normal((65536,))
r = mx.random.normal((65536,))
mx.eval(x, r)

for s in (0.3, 1 / 3, 128 ** -0.5, 0.7071067811865476):
    eager = (x * r) * s
    compiled = mx.compile(lambda x, r, s=s: (x * r) * s)(x, r)
    mx.eval(eager, compiled)
    print(f"s={s!r}: {int((eager != compiled).sum())} of 65536 elements differ")

# same constant passed as an input instead of captured: identical
sa = mx.array(128 ** -0.5, mx.float32)
eager = (x * r) * (128 ** -0.5)
compiled = mx.compile(lambda x, r, sa: (x * r) * sa)(x, r, sa)
mx.eval(eager, compiled)
print("as input:", int((eager != compiled).sum()))

Output on mlx 0.32.2:

s=0.3: 0 of 65536 elements differ
s=0.3333333333333333: 65536 of 65536 elements differ
s=0.08838834764831845: 60474 of 65536 elements differ
s=0.7071067811865476: 60474 of 65536 elements differ
as input: 0

Expected behavior

A compiled function should compute the same thing as its eager form; a captured constant should be embedded exactly. Using std::numeric_limits<T>::max_digits10 in print_float_constant (or emitting std::hexfloat) makes the embedded literal round-trip.

Desktop

  • OS: macOS 26.4
  • Chip: Apple M3 Ultra
  • Python 3.13, mlx 0.32.2 (also present in main's compiled.h as of 2026-09-15)

Additional context

Found while compiling a linear-attention decode step (x * rsqrt(sum(x*x) + eps) * head_dim**-0.5) and checking it bit-for-bit against eager: after accounting for mx.sigmoid's fast exp under the JIT, this was the last remaining difference. Passing the scale as an array input works around it.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start in mlx/backend/common/compiled.h at print_float_constant and run the provided Python reproduction to observe captured constants differing from eager results. Check both float32 and double formatting paths; done means compiled functions preserve captured scalar values closely enough to match eager computation, including the examples that currently differ.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp, python
Domain
backend, compilers
Issue type
Bug
Difficulty
2/5
Estimated time
1-3 hours
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
85/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.