Fix unit_diagonal and non-contiguous reads in MLX linalg

pytensor/link/mlx/dispatch/basic.py

@@ -132,7 +132,9 @@ def mlx_typify_tensor(data, dtype=None, **kwargs):
     # and it does so on the CPU too, where float64 is perfectly usable
     if dtype is None and data.dtype == np.float64 and float64_supported():
         dtype = mx.float64
-    return _nan_safe_constant(data, dtype=dtype)
+    # MLX's elementwise kernels misread a non-contiguous buffer, and rewriting
+    # produces such arrays -- `triu` of a transpose, for one
+    return _nan_safe_constant(np.asarray(data, order="C"), dtype=dtype)
 
 
 @mlx_typify.register(slice)

pytensor/link/mlx/dispatch/linalg/solvers.py

@@ -33,12 +33,22 @@ def solve(a, b):
 @mlx_funcify.register(SolveTriangular)
 def mlx_funcify_SolveTriangular(op, node, **kwargs):
     lower = op.lower
+    unit_diagonal = op.unit_diagonal
     A_dtype = getattr(mx, node.inputs[0].dtype)
     b_dtype = getattr(mx, node.inputs[1].dtype)
 
     def solve_triangular(A, b):
+        A = A.astype(stream=mx.cpu, dtype=A_dtype)
+
+        if unit_diagonal:
+            # MLX's `solve_triangular` has no `unit_diagonal`. LAPACK's `trtrs`
+            # never reads the diagonal in that mode, so overwriting it with ones
+            # gives the same answer.
+            diagonal_mask = mx.eye(A.shape[-1], dtype=mx.bool_, stream=mx.cpu)
+            A = mx.where(diagonal_mask, mx.array(1, dtype=A_dtype), A, stream=mx.cpu)
+
         return mx.linalg.solve_triangular(
-            A.astype(stream=mx.cpu, dtype=A_dtype),
+            A,
             b.astype(stream=mx.cpu, dtype=b_dtype),
             upper=not lower,
             stream=mx.cpu,

tests/link/mlx/linalg/test_decomposition.py

@@ -5,7 +5,7 @@
 from packaging.version import parse as V
 
 import pytensor.tensor as pt
-from pytensor import config
+from pytensor import config, grad
 from pytensor.tensor.linalg.decomposition import lu, svd
 from pytensor.tensor.linalg.decomposition.cholesky import cholesky
 from tests.link.mlx.test_basic import compare_mlx_and_py, mlx_mode
@@ -131,6 +131,27 @@ def test_mlx_lu_factor():
     compare_mlx_and_py([A], out, [A_val])
 
 
+def test_mlx_lu_solve_grad():
+    """The pullback folds non-contiguous constants, which MLX must read correctly."""
+    rng = np.random.default_rng(15)
+
+    A = pt.tensor(name="A", shape=(5, 5))
+    b = pt.tensor(name="b", shape=(5,))
+    A_val = rng.normal(size=(5, 5)).astype(config.floatX)
+    b_val = rng.normal(size=(5,)).astype(config.floatX)
+
+    x = pt.linalg.lu_solve(pt.linalg.lu_factor(A), b)
+
+    # `fast_run` so the permutation `arange` in the pullback folds to a constant,
+    # which MLX's `arange` requires
+    compare_mlx_and_py(
+        [A, b],
+        [grad(x.sum(), A)],
+        [A_val, b_val],
+        mlx_mode=mlx_mode.including("fast_run"),
+    )
+
+
 def test_mlx_pivot_to_permutations():
     rng = np.random.default_rng(15)
 

tests/link/mlx/linalg/test_solvers.py

@@ -44,22 +44,27 @@ def test_mlx_solve(assume_a):
         )
 
 
+@pytest.mark.parametrize(
+    "unit_diagonal", [False, True], ids=["full_diagonal", "unit_diagonal"]
+)
 @pytest.mark.parametrize("lower", [True, False], ids=["lower", "upper"])
-def test_mlx_SolveTriangular(lower):
+def test_mlx_SolveTriangular(lower, unit_diagonal):
     rng = np.random.default_rng(15)
 
     A = pt.tensor("A", shape=(5, 5))
     b = pt.tensor("B", shape=(5, 5))
 
+    # A diagonal far from one, so ignoring `unit_diagonal` gives a different answer
     A_val = rng.normal(size=(5, 5)).astype(config.floatX)
+    A_val[np.diag_indices(5)] = rng.uniform(3, 4, size=5).astype(config.floatX)
     b_val = rng.normal(size=(5, 5)).astype(config.floatX)
 
     out = pt.linalg.solve_triangular(
         A,
         b,
         trans=0,
         lower=lower,
-        unit_diagonal=False,
+        unit_diagonal=unit_diagonal,
     )
     compare_mlx_and_py(
         [A, b],

tests/link/mlx/test_basic.py

@@ -23,7 +23,7 @@
 
 
 mx = pytest.importorskip("mlx.core")
-from pytensor.link.mlx.dispatch.basic import convert_dtype_to_mlx
+from pytensor.link.mlx.dispatch.basic import convert_dtype_to_mlx, mlx_typify
 
 
 optimizer = RewriteDatabaseQuery(include=["mlx"], exclude=MLX._optimizer.exclude)
@@ -355,3 +355,25 @@ def test_nan_array_constant():
     compare_mlx_and_py(
         [x], [x + c], [np.array([10.0, 20.0, 30.0], dtype=config.floatX)]
     )
+
+
+@pytest.mark.parametrize(
+    "data",
+    [
+        np.asfortranarray(np.arange(6, dtype="float32").reshape(2, 3)),
+        np.arange(6, dtype="float32").reshape(2, 3).T,
+        np.arange(6, dtype="float32").reshape(2, 3),
+        np.array(2.0, dtype="float32"),
+    ],
+    ids=["fortran_order", "transposed_view", "already_contiguous", "scalar"],
+)
+def test_mlx_typify_row_contiguous(data):
+    # MLX's elementwise kernels misread a non-contiguous buffer, and rewriting
+    # produces such arrays (``triu`` of a transpose, for one). Only the full LU
+    # pullback reproduces the wrong answer, so the invariant is pinned here.
+    result = mlx_typify(data)
+    as_numpy = np.asarray(result)
+
+    np.testing.assert_array_equal(as_numpy, data)
+    assert as_numpy.flags["C_CONTIGUOUS"]
+    assert result.shape == data.shape