mlx
mlx copied to clipboard
[BUG] `mx.linalg.cholesky` is less precise than `scipy.linalg.cholesky`?
Describe the bug
For either float32 or float64 inputs, mx.linalg.cholesky seems to occasionally disagree with the LAPACK implementation used by scipy. Highly possible I'm missing about doing full-precision computation -- apologies in advance if that's the case. I checked the and didn't see any resources about higher precision. The problem seems to exist at half precision though.
To Reproduce
Raises:
import mlx.core as mx
import numpy as np
from scipy import linalg
for floatX in ['float32', 'float64']:
res = np.empty(1000, dtype=bool)
mx_type = getattr(mx, floatX)
for i in range(1000):
A_val = rng.normal(size=(5, 5)).astype(floatX)
A_val = A_val @ A_val.T
b_val = rng.normal(size=(5, 5)).astype(floatX)
A_val_mx = mx.array(A_val).astype(mx_type, stream=mx.cpu)
chol = mx.linalg.cholesky(A_val_mx, upper=True, stream=mx.cpu).astype(mx_type, stream=mx.cpu)
res[i] = np.allclose(chol, linalg.cholesky(A_val))
print(f'{floatX} agreement: {res.mean():0.2%}')
Output:
float32 agreement: 88.10%
float64 agreement: 88.00%
Expected behavior
MLX agrees with LAPACK implementation to high precision (np.allclose defaults are rtol=1e-5, atol=1e-8)
Desktop (please complete the following information):
- OS Version: MacOS 15.4 (24E248)
- Version 0.25.0
Additional context None