mx.random.normal: bit-exact divergence on Apple M1 Max vs M3 Ultra / M5

Description

mx.random.normal(shape=..., key=mx.random.key(seed)) produces different float32 outputs on Apple M1 Max vs Apple M3 Ultra and Apple M5, given identical MLX version, identical Python, and identical inputs. mx.random.uniform and mx.random.split are unaffected — both produce bit-identical output across all three machines.

This breaks bit-exact reproducibility for any code that relies on mx.random.normal with a directly-constructed key, across Apple Silicon hardware generations.

Reproducer

import mlx.core as mx
import numpy as np
import hashlib

key = mx.random.key(0)
a = mx.random.normal(shape=(8,), key=key)
mx.eval(a)
print(hashlib.sha256(np.asarray(a, dtype=np.float32).tobytes()).hexdigest())

Observed output (MLX 0.31.1, Python 3.14.3-4, macOS 26)

Chip	`sha256` of `normal(shape=(8,), key=mx.random.key(0))`
Apple M5	`a662607828c52cafea1f79d98f09bd81855784116acd23364a6130eb037d6a5e`
Apple M3 Ultra	`a662607828c52cafea1f79d98f09bd81855784116acd23364a6130eb037d6a5e`
Apple M1 Max	`8a8c5a6737feea9eeee3c7056f4699250e6b8440db8930fd10495ada9117c05a`

Same pattern for shape=(8,) key=mx.random.key(42) and shape=(16, 4) key=mx.random.key(0) — M1 Max always differs from {M5, M3 Ultra}.

Adjacent observations — narrowing the cause

The following calls produce bit-identical output across all three machines:

Probe	M5	M3 Ultra	M1 Max	All match ?
`mx.random.key(0)` materialised	✓	✓	✓	yes
`mx.random.split(mx.random.key(7))` (both halves)	✓	✓	✓	yes
`mx.random.uniform(low=-1, high=1, shape=(8,), key=mx.random.key(0))`	✓	✓	✓	yes
`mx.random.normal(shape=(4,), key=<split-child-of-key(7)>)`	✓	✓	✓	yes
`mx.random.normal(shape=(8,), key=mx.random.key(0))`	✓	✓	❌	no
`mx.random.normal(shape=(8,), key=mx.random.key(42))`	✓	✓	❌	no
`mx.random.normal(shape=(16, 4), key=mx.random.key(0))`	✓	✓	❌	no

So the divergence is specifically in mx.random.normal on Apple M1 Max when the key is the raw output of mx.random.key(seed). The key tensor itself is identical; mx.random.split returns identical halves; mx.random.uniform produces identical output with the same raw key. Only mx.random.normal with the raw key diverges, only on M1 Max.

It is plausible (but not confirmed) that mx.random.normal with a split-derived key matches across all three, but mx.random.normal with a directly-constructed key takes a different kernel path on M1 Max.

Environment

All three machines:

MLX 0.31.1
Python 3.14.3 (M5) / 3.14.4 (M3 Ultra, M1 Max)
macOS 26.x (Tahoe), arm64
Same uv lockfile, single MLX wheel

Why this matters (downstream)

We hit this in a research project (dream-of-kiki) that runs an R1 bit-exact reproducibility suite. After shipping new ops that call mx.random.normal with seeded keys, the hashes diverge between M1 Max and M5/M3 Ultra even at the same commit, breaking cross-machine equivalence that previously held for ops that didn't touch normal. Full milestone: docs/milestones/r1-cross-machine-m5-vs-m1-2026-05-20.md.

We've put a golden_hashes.json per-hardware-family workaround in place locally, but it would be much cleaner if mx.random.normal produced the same bytes across Apple Silicon generations.

Full repro script

A self-contained script that emits a JSON dump of 8 different mx.random.* probes (including the matching ones above as controls): [/tmp/mlx_rng_repro.py — paste here on request].

Question

Is this an intentional kernel divergence (different SIMD path / Philox internal variant per hardware family) or a regression worth fixing?
If intentional, is there a recommended way to opt into a portable code path for mx.random.normal?

milestone: R1 probe 3-machines + MLX upstream

feat(r1): split golden_hashes per chip family

docs: qualify R1 within-machine vs cross-machine

Forensic update — likely root cause is `erfinv` FMA chain

Tracing the call path in mlx/random.cpp (v0.31.2, commit 68cf2fd):

// mlx/random.cpp:174-204 — mx.random.normal implementation
array normal(...) {
  auto low  = array(std::nextafter(-1.0f, 0.0f), float32);
  auto high = array(1.0f, float32);
  auto samples = uniform(low, high, shape, float32, key, stream);   // (1)
  auto applied_scale = array(std::sqrt(2.0), dtype);
  samples = astype(erfinv(samples, stream), dtype, stream);          // (2)
  samples = multiply(applied_scale, samples, stream);                // (3)
  return samples;
}

The probe in the original issue isolates the divergence to step (2) erfinv:

Step (1) — uniform(low=-1+eps, high=1, key=mx.random.key(0)) matches across M5, M3 Ultra, M1 Max (verified by mx.random.uniform(low=-1.0, high=1.0, ...) probe).
Step (3) — multiply(sqrt(2), ...) is a pure element-wise op and would not be hardware-dependent for finite inputs.

That leaves step (2) — the erfinv Metal kernel in mlx/backend/metal/kernels/erf.h:42. Its body is a 10-step Horner-form polynomial entirely built from metal::fma calls:

// erf.h:42-69 — abridged
float erfinv(float a) {
  auto t = metal::fma(a, 0.0f - a, 1.0f);
  t = metal::log(t);
  float p;
  if (metal::abs(t) > 6.125f) {
    p = 3.03697567e-10f;
    p = metal::fma(p, t, 2.93243101e-8f);
    p = metal::fma(p, t, 1.22150334e-6f);
    // ... 7 more chained metal::fma steps ...
    p = metal::fma(p, t, 8.40016484e-1f);
  } else {
    // symmetric 10-step chain
  }
  return a * p;
}

metal::fma(a, b, c) = a*b + c with a single IEEE-754 rounding (intermediate result is held in extended precision). If the Metal compiler lowers metal::fma to a hardware FMA on M5 / M3 Ultra but to a separate multiply-then-add path with two roundings on M1 Max (or vice versa), the entire polynomial cascade diverges at float32 precision.

The 10-step chain compounds the per-step ulp differences: each step's r1 and r3 reduction stages can differ in low bits, and by the time we exit the loop we have float32 outputs that bit-differ in the last few mantissa bits across hardware.

Why some normal calls match and others don't

In the original probe, mx.random.normal(shape=(4,), key=split_derived_key) matched all three machines while mx.random.normal(shape=(8,), key=raw_mx.random.key(0)) diverged on M1 Max. Hypothesis: the Metal kernel launch grid for the polynomial differs by shape (threadgroup size, vectorization) and only some launch shapes route through the FMA-divergent path on M1 Max. A confirming test would be:

for shape in [(1,), (2,), (4,), (8,), (16,), (32,), (64,)]:
    a = mx.random.normal(shape=shape, key=mx.random.key(0))
    mx.eval(a)
    print(shape, hashlib.sha256(np.asarray(a, np.float32).tobytes()).hexdigest()[:16])

…run on each machine and diff. If only specific shapes diverge, the launch-grid/SIMD-width hypothesis is confirmed.

C++ unit test demonstrating the divergence

A minimal unit test that could live in tests/random_tests.cpp:

TEST_CASE("random.normal cross-machine determinism") {
  // Same input across machines — this hash MUST match
  // on every Apple Silicon family for the R1 contract.
  using namespace mlx::core;
  auto key = random::key(0);
  auto out = random::normal({8}, float32, std::nullopt, std::nullopt, key);
  out.eval();
  
  // Pre-recorded expected bytes (captured on a reference machine,
  // would need to be promoted to a CI-pinned baseline):
  const std::vector<uint8_t> expected_bytes = {
    0x1e, 0x9e, 0xa5, 0x3d, 0x28, 0xc2, 0xc5, 0xbe,
    0xea, 0x55, 0xc0, 0xbe, 0xf5, 0xa0, 0xd5, 0x3f,
    0x0b, 0x4e, 0xa3, 0xbf, 0xfa, 0xa0, 0x07, 0x40,
    0xc6, 0xb3, 0x5b, 0xbf, 0x49, 0xb7, 0x90, 0x3f,
  };
  
  const auto out_bytes = std::vector<uint8_t>(
    static_cast<const uint8_t*>(out.data<float>()),
    static_cast<const uint8_t*>(out.data<float>()) + 32);
  
  CHECK(out_bytes == expected_bytes);  // Fails on M1 Max only
}

The first 32 bytes shown are M5 / M3 Ultra; M1 Max produces different bytes.

Possible mitigations (sorted by intrusiveness)

Determinism-preserving polynomial. Replace the chained metal::fma calls in erfinv with explicit multiply-then-add (p = p * t; p = p + coeff). Costs ~2× in throughput on the polynomial — but erfinv is rarely the bottleneck. Trade perf for cross-hardware reproducibility.
CPU fallback for erfinv. Route the erfinv op to the CPU backend when a deterministic-across-hardware mode is requested. Requires an opt-in flag (mx.set_deterministic(True) or similar) — non-trivial API surface.
Document the constraint. The simplest fix: document that mx.random.normal bit-exactness is per-chip-family-conditional, and users requiring cross-hardware bit-exactness should pre-compute on one machine and ship the outputs. Zero code change.

I don't have strong opinions about which path the project should take — the research project that surfaced this hit it accidentally and has migrated to per-chip-family golden hashes locally (see hypneum-lab/dream-of-kiki commit deeb1e8). Documentation alone would have been sufficient for us; the FMA forensic is shared in case it's useful to the maintainers.

Happy to test patches on M5 / M3 Ultra / M1 Max if a candidate fix lands — full per-machine test infrastructure is already in place at hypneum-lab/dream-of-kiki/tests/reproducibility/.

Empirical update — added 4th host (M3 Pro, `applegpu_g15s`) confirms the `g13 vs g15+` frontier

Following up the original issue with a 4-host triangulation at fixed MLX 0.31.2 (rules out version regression) and isolated applegpu architecture strings reported by mx.metal.device_info():

Host	Chip	`applegpu` arch	`device_info()` reports
macm1	Apple M1 Max	`applegpu_g13s`	`{'device_name': 'Apple M1 Max', 'architecture': 'applegpu_g13s', ...}`
macM3	Apple M3 Pro	`applegpu_g15s`	`{'device_name': 'Apple M3 Pro', 'architecture': 'applegpu_g15s', ...}`
studio	Apple M3 Ultra	`applegpu_g15d`	`{'device_name': 'Apple M3 Ultra', 'architecture': 'applegpu_g15d', ...}`
grosmac	Apple M5	`applegpu_g17g`	`{'device_name': 'Apple M5', 'architecture': 'applegpu_g17g', ...}`

Cross-arch sha256 table (MLX 0.31.2 everywhere)

Tested across 3 sample sizes × 5 seeds × normal and uniform. Hashes are sha256 hex over f"{v:.8e}" per-value, truncated to 24 chars:

`mx.random.normal` (8 tests × 4 hosts)

Test	M1 Max (g13s)	M3 Pro (g15s)	M3 Ultra (g15d)	M5 (g17g)
`normal(seed=0, n=1000)`	`7530424b...`	`b628fb71...`	`b628fb71...`	`b628fb71...`
`normal(seed=0, n=10000)`	`be71f90a...`	`ef71df1d...`	`ef71df1d...`	`ef71df1d...`
`normal(seed=0, n=100000)`	`b488ea90...`	`750379c9...`	`750379c9...`	`750379c9...`
`normal(seed=1, n=10000)`	`8bd6ad78...`	`8e047aef...`	`8e047aef...`	`8e047aef...`
`normal(seed=42, n=10000)`	`372f5de0...`	`17767b0d...`	`17767b0d...`	`17767b0d...`
`normal(seed=100, n=10000)`	`5cb3349a...`	`c54bcfb3...`	`c54bcfb3...`	`c54bcfb3...`
`normal(seed=1234, n=10000)`	`1c0f6499...`	`285a378d...`	`285a378d...`	`285a378d...`

`mx.random.uniform` control (2 tests × 4 hosts) — bit-exact everywhere

Test	M1 Max	M3 Pro	M3 Ultra	M5
`uniform(seed=0, n=10000)`	`af192db2...`	`af192db2...`	`af192db2...`	`af192db2...`
`uniform(seed=0, n=100000)`	`0db1b61b...`	`0db1b61b...`	`0db1b61b...`	`0db1b61b...`

Frontière empirique nette

8 normal tests × 4 hosts = 32 hashes. Distribution :

M1 Max (g13s) produces its own unique sha256 on all 8 tests — distinct from every other host.
M3 Pro (g15s), M3 Ultra (g15d), M5 (g17g) are mutually bit-exact on all 8 tests despite having three different applegpu arch strings (g15s ≠ g15d ≠ g17g).
uniform is bit-exact on all 4 hosts.

The divergence is therefore localized to uniform → erfinv transformation (cf. previous comment on mlx/backend/metal/kernels/erf.h:42's 10-step Horner polynomial), and the frontier is g13 (M1 family, Apple7) vs g15+ (Apple8 onwards). The Metal compiler appears to emit identical kernel bytecode for erfinv from g15 onward, but a different FMA chain layout for g13.

Mechanistic hypothesis (consistent with #2205)

This matches angeloskath's earlier explanation in #2205 ("the grid needs to be a different size for the M1 vs the M3 basically this means that you can't run 1024 threads per thread group on the M1"). The M1's tighter register file (Rosenzweig's M1 GPU reverse-engineering documents ~256 half-word registers per thread vs. larger budget on g15+) forces the Metal compiler to choose a different SIMD-width / FMA-fusion strategy for erfinv's polynomial body — same source code, different bytecode emitted.

Reproducer (5 lines, runs in any MLX install)

import hashlib
import mlx.core as mx
mx.random.seed(0)
x = mx.random.normal((10000,))
mx.eval(x)
import numpy as np
arr = np.asarray(x)
print(f"{mx.metal.device_info()['architecture']}: "
      f"{hashlib.sha256(','.join(f'{v:.8e}' for v in arr).encode()).hexdigest()[:24]}")

Expected output on Apple M1 family: applegpu_g13s: be71f90acef4a52e36ca89ad
Expected output on Apple M3+ : applegpu_g15s/g15d/g17g: ef71df1dfcccab1cb156c6fc

What this changes for the documentation request

The original ask remains: document this in mx.random.normal's docstring. The added evidence shows:

It's not a version regression — same MLX 0.31.2 wheel everywhere.
It's not a M1-Max-specific quirk — it's the entire g13 (Apple7) family vs everything Apple8+.
It's specifically erfinv, not the whole normal pipeline — uniform and split are unaffected.
The behaviour is stable across MLX 0.31.2 → 0.32.0.dev (verified separately on M1 Max).

A minimal docstring note like "Note: mx.random.normal is not bit-exact across Apple Silicon GPU architectures (g13 vs g15+) — use a fixed-architecture worker for reproducibility-critical workflows" would close the gap.

🔬 4-host triangulation by @electron-rare (M1 Max, M3 Pro, M3 Ultra, M5) at MLX 0.31.2. Generated as part of the nerve-wml gap-analysis remediation.

Re-test at MLX 0.32.4: M3 Pro, M3 Ultra and M5 still agree; M1 Max not re-run

Re-ran the probes at 0.32.4 on M3 Pro (g15s), M3 Ultra (g15d) and M5 (g17g); all 19 are bit-identical across the three. normal(shape=(8,), key=mx.random.key(0)) gives a662607828c52cafea1f79d9, the same value reported for M5 / M3 Ultra in the issue body at 0.31.1. I did not re-run the M1 Max (g13s), so this says nothing new about the g13 divergence.

Hashes: first 24 hex chars of sha256 over the raw float32 bytes, as in the issue body. My May 20 comment used mx.random.seed() and hashed formatted strings, so its values are not comparable.

Probe	0.32.4 (g15s = g15d = g17g)
`normal((8,), key(0))`	`a662607828c52cafea1f79d9`
`normal((8,), key(42))`	`7ef9a0d55131a5b50cb4d5b5`
`normal((16, 4), key(0))`	`b2b3b19fe3512c1b88f78032`
`normal((4,), split(key(7))[0])`	`a42df6b2e9a815bd93b05c59`
`normal((1000,), key(0))`	`46b0d4be6635990603a76631`
`normal((10000,), key(0))`	`03ef4378f4c0325a232d5edd`
`normal((100000,), key(0))`	`72e6ab35a28a4b2f8290d4e8`
`normal((10000,), key(1) / key(42) / key(100) / key(1234))`	`39364402dbec903e7bc5e5f3` / `5fe34c2a1665002b0b378b97` / `b96f53140ef1f69ec89b467f` / `141599632fb7a7161ce5ec15`
`uniform(-1, 1, (8,), key(0))` (control)	`f8646c5b74ac4c237e50710a`
`uniform((10000,), key(0))` (control)	`386c7b9e80903e37662a3836`
`uniform((100000,), key(0))` (control)	`fe91dc3993634df1097d92ed`
`erfinv(linspace(-0.999, 0.999, 10001))` (new)	`a983396ba86f1a158e22aca3`
`erf(linspace(-0.999, 0.999, 10001))` (new)	`d3c14ad20901a5fbe6495f71`

The last two rows apply erfinv and erf to a fixed input, without the RNG; they have not been run on g13.

What I checked

Release notes v0.32.0 to v0.32.4: the only erfinv entry is "Reject complex in erf and erfinv" (Reject complex in erf and erfinv #4243). No random-related entry (test: Upcast the random numbers before computing their average #3488, Fix captured random state in compile #3828, Reject negative dimensions in broadcast_to and random shapes #4046, Fix random distribution dtype error messages #4131) describes a numerical change.
random.cpp at v0.32.4: normal (l.179) still draws uniform(nextafter(-1, 0), 1) (l.194-196), applies erfinv (l.202), multiplies by sqrt(2) (l.203).
kernels/erf.h at v0.32.4: erfinv (l.42-69) is metal::log (l.44) then a metal::fma polynomial; kernels are built with -fno-fast-math (kernels/CMakeLists.txt l.18). I have not inspected per-architecture compiler output, so I cannot say whether log or the FMA chain differs on g13.

Question for the maintainers. Would either of these be acceptable?
(a) A portable erfinv path for random.normal, accepting a throughput cost and a likely one-time change of normal outputs.
(b) A docstring note that mx.random.normal is not guaranteed bit-exact across GPU architectures (observed at 0.31.x: g13 vs g15+).

I can open a PR for (b). Probe script below.

Environment: MLX 0.32.4 via pip --require-hashes; macOS 27.0.1 / 27.2; Python 3.12.15 (M5), 3.13.9 (M3 Pro), 3.13.16 (M3 Ultra).

probe.py

"""Re-test of ml-explore/mlx#3568: cross-architecture bit-exactness of
mx.random.normal vs mx.random.uniform. Prints one JSON object."""
import hashlib, json, platform, subprocess, sys
import mlx.core as mx
import numpy as np

def h(a):
    mx.eval(a)
    return hashlib.sha256(np.asarray(a, dtype=np.float32).tobytes()).hexdigest()[:24]

probes = {}
probes["key0_materialised"] = h(mx.random.key(0).astype(mx.float32))
k7a, k7b = mx.random.split(mx.random.key(7))
probes["split_key7_a"] = h(k7a.astype(mx.float32)); probes["split_key7_b"] = h(k7b.astype(mx.float32))
probes["uniform_key0_8"] = h(mx.random.uniform(low=-1, high=1, shape=(8,), key=mx.random.key(0)))
probes["normal_splitchild_key7_4"] = h(mx.random.normal(shape=(4,), key=k7a))
probes["normal_key0_8"] = h(mx.random.normal(shape=(8,), key=mx.random.key(0)))
probes["normal_key42_8"] = h(mx.random.normal(shape=(8,), key=mx.random.key(42)))
probes["normal_key0_16x4"] = h(mx.random.normal(shape=(16, 4), key=mx.random.key(0)))
for seed, n in [(0, 1000), (0, 10000), (0, 100000), (1, 10000), (42, 10000), (100, 10000), (1234, 10000)]:
    probes[f"normal_seed{seed}_n{n}"] = h(mx.random.normal(shape=(n,), key=mx.random.key(seed)))
for n in (10000, 100000):
    probes[f"uniform_seed0_n{n}"] = h(mx.random.uniform(shape=(n,), key=mx.random.key(0)))
# erfinv alone on a fixed deterministic input (the suspected non-portable op)
x = mx.linspace(-0.999, 0.999, 10001, dtype=mx.float32)
probes["erfinv_linspace_10001"] = h(mx.erfinv(x))
probes["erf_linspace_10001"] = h(mx.erf(x))
info = {}
try:
    info = mx.metal.device_info()
except Exception as e:
    info = {"error": str(e)}
chip = subprocess.run(["sysctl", "-n", "machdep.cpu.brand_string"], capture_output=True, text=True).stdout.strip()
print(json.dumps({"host": platform.node(), "chip": chip, "macos": platform.mac_ver()[0],
                  "python": sys.version.split()[0], "mlx": mx.__version__,
                  "device_info": {k: (v if isinstance(v, (str, int, float)) else str(v)) for k, v in info.items()},
                  "probes": probes}, indent=1))

docs: note random.normal cross-GPU bit-exactness

Navigation Menu

Repository navigation

mx.random.normal: bit-exact divergence on Apple M1 Max vs M3 Ultra / M5 #3568

Description

Reproducer

Observed output (MLX 0.31.1, Python 3.14.3-4, macOS 26)

Adjacent observations — narrowing the cause

Environment

Why this matters (downstream)

Full repro script

Question

Forensic update — likely root cause is `erfinv` FMA chain

Why some normal calls match and others don't

C++ unit test demonstrating the divergence

Possible mitigations (sorted by intrusiveness)

Empirical update — added 4th host (M3 Pro, `applegpu_g15s`) confirms the `g13 vs g15+` frontier

Cross-arch sha256 table (MLX 0.31.2 everywhere)

`mx.random.normal` (8 tests × 4 hosts)

`mx.random.uniform` control (2 tests × 4 hosts) — bit-exact everywhere

Frontière empirique nette

Mechanistic hypothesis (consistent with #2205)

Reproducer (5 lines, runs in any MLX install)

What this changes for the documentation request

Re-test at MLX 0.32.4: M3 Pro, M3 Ultra and M5 still agree; M1 Max not re-run

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

mx.random.normal: bit-exact divergence on Apple M1 Max vs M3 Ultra / M5 #3568

Description

Description

Reproducer

Observed output (MLX 0.31.1, Python 3.14.3-4, macOS 26)

Adjacent observations — narrowing the cause

Environment

Why this matters (downstream)

Full repro script

Question

Activity

electron-rare commented on May 20, 2026

Forensic update — likely root cause is erfinv FMA chain

Why some normal calls match and others don't

C++ unit test demonstrating the divergence

Possible mitigations (sorted by intrusiveness)

electron-rare commented on May 20, 2026

Empirical update — added 4th host (M3 Pro, applegpu_g15s) confirms the g13 vs g15+ frontier

Cross-arch sha256 table (MLX 0.31.2 everywhere)

mx.random.normal (8 tests × 4 hosts)

mx.random.uniform control (2 tests × 4 hosts) — bit-exact everywhere

Frontière empirique nette

Mechanistic hypothesis (consistent with #2205)

Reproducer (5 lines, runs in any MLX install)

What this changes for the documentation request

electron-rare commented on Oct 31, 2026

Re-test at MLX 0.32.4: M3 Pro, M3 Ultra and M5 still agree; M1 Max not re-run

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions

Forensic update — likely root cause is `erfinv` FMA chain

Empirical update — added 4th host (M3 Pro, `applegpu_g15s`) confirms the `g13 vs g15+` frontier

`mx.random.normal` (8 tests × 4 hosts)

`mx.random.uniform` control (2 tests × 4 hosts) — bit-exact everywhere