mlx icon indicating copy to clipboard operation
mlx copied to clipboard

[BUG] `mx.linalg.cholesky` is less precise than `scipy.linalg.cholesky`?

Open jessegrabowski opened this issue 11 months ago • 0 comments

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

jessegrabowski avatar Apr 19 '25 22:04 jessegrabowski