diff --git a/pytensor/link/mlx/dispatch/linalg/solvers.py b/pytensor/link/mlx/dispatch/linalg/solvers.py index 12be312680..2f3f31416a 100644 --- a/pytensor/link/mlx/dispatch/linalg/solvers.py +++ b/pytensor/link/mlx/dispatch/linalg/solvers.py @@ -4,6 +4,7 @@ from pytensor.link.mlx.dispatch.basic import mlx_funcify from pytensor.tensor.linalg.solvers.general import Solve +from pytensor.tensor.linalg.solvers.psd import CholeskySolve from pytensor.tensor.linalg.solvers.triangular import SolveTriangular @@ -44,3 +45,22 @@ def solve_triangular(A, b): ) return solve_triangular + + +@mlx_funcify.register(CholeskySolve) +def mlx_funcify_CholeskySolve(op, node, **kwargs): + lower = op.lower + c_dtype = getattr(mx, node.inputs[0].dtype) + b_dtype = getattr(mx, node.inputs[1].dtype) + + # MLX has no cho_solve, so with A = L L.T we solve L y = b then L.T x = y. + def cho_solve(c, b): + c = c.astype(stream=mx.cpu, dtype=c_dtype) + b = b.astype(stream=mx.cpu, dtype=b_dtype) + c_T = mx.swapaxes(c, -1, -2, stream=mx.cpu) + L, L_T = (c, c_T) if lower else (c_T, c) + + y = mx.linalg.solve_triangular(L, b, upper=False, stream=mx.cpu) + return mx.linalg.solve_triangular(L_T, y, upper=True, stream=mx.cpu) + + return cho_solve diff --git a/tests/link/mlx/linalg/test_solvers.py b/tests/link/mlx/linalg/test_solvers.py index d70c8bd0a4..0028296450 100644 --- a/tests/link/mlx/linalg/test_solvers.py +++ b/tests/link/mlx/linalg/test_solvers.py @@ -44,8 +44,8 @@ def test_mlx_solve(assume_a): ) -@pytest.mark.parametrize("lower, trans", [(False, False), (True, True)]) -def test_mlx_SolveTriangular(lower, trans): +@pytest.mark.parametrize("lower", [True, False], ids=["lower", "upper"]) +def test_mlx_SolveTriangular(lower): rng = np.random.default_rng(15) A = pt.tensor("A", shape=(5, 5)) @@ -70,3 +70,61 @@ def test_mlx_SolveTriangular(lower, trans): np.testing.assert_allclose, atol=1e-6, rtol=1e-6, strict=True ), ) + + +@pytest.mark.parametrize("batch_shape", [(), (3,)], ids=["core", "batched"]) +@pytest.mark.parametrize("lower", [True, False], ids=["lower", "upper"]) +@pytest.mark.parametrize("b_ndim", [1, 2], ids=["b_vec", "b_mat"]) +def test_mlx_CholeskySolve(batch_shape, lower, b_ndim): + rng = np.random.default_rng(15) + n = 5 + b_shape = (*batch_shape, n) if b_ndim == 1 else (*batch_shape, n, 3) + + C = pt.tensor("C", shape=(*batch_shape, n, n)) + b = pt.tensor("b", shape=b_shape) + + out = pt.linalg.cho_solve((C, lower), b, b_ndim=b_ndim) + + A_val = rng.normal(size=(*batch_shape, n, n)).astype(config.floatX) + A_val = A_val @ np.swapaxes(A_val, -1, -2) + n * np.eye(n, dtype=config.floatX) + C_val = np.linalg.cholesky(A_val) + if not lower: + C_val = np.swapaxes(C_val, -1, -2).copy() + + b_val = rng.normal(size=b_shape).astype(config.floatX) + + compare_mlx_and_py( + [C, b], + [out], + [C_val, b_val], + mlx_mode=mlx_mode, + assert_fn=partial( + np.testing.assert_allclose, atol=1e-6, rtol=1e-6, strict=True + ), + ) + + +def test_mlx_CholeskySolve_mixed_dtypes(): + rng = np.random.default_rng(15) + n = 5 + + C = pt.tensor("C", shape=(n, n), dtype="float32") + b = pt.tensor("b", shape=(n,), dtype="float64") + + out = pt.linalg.cho_solve((C, True), b, b_ndim=1) + assert out.type.dtype == "float64" + + A_val = rng.normal(size=(n, n)) + A_val = A_val @ A_val.T + n * np.eye(n) + C_val = np.linalg.cholesky(A_val).astype("float32") + b_val = rng.normal(size=(n,)) + + compare_mlx_and_py( + [C, b], + [out], + [C_val, b_val], + mlx_mode=mlx_mode, + assert_fn=partial( + np.testing.assert_allclose, atol=1e-5, rtol=1e-5, strict=True + ), + )