mx.compile inlines float scalar constants with 7 significant digits, so compiled results differ from eager by 1 ulp
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'scompiled.has 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
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- 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