MLX backend silently ignores unit_diagonal in SolveTriangular

[jessegrabowski]

mlx_funcify_SolveTriangular reads op.lower but never op.unit_diagonal, and mx.linalg.solve_triangular has no such argument, so the flag is dropped and the stored diagonal gets used instead of ones. Nothing raises. It also breaks lu_solve, whose L is unit-triangular packed into the LU array with U's diagonal in those slots, so any graph where reuse_decomposition_multiple_solves collapses two solves onto one factorization silently returns wrong numbers on MLX.

import numpy as np
import pytensor
import pytensor.tensor as pt

A = pt.matrix("A", dtype="float32")
b = pt.vector("b", dtype="float32")
out = pt.linalg.solve_triangular(A, b, lower=True, unit_diagonal=True, b_ndim=1)

# unit_diagonal=True must ignore the stored diagonal, i.e. solve
# [[1,0,0],[1,1,0],[1,1,1]] x = b  ->  x = [1, 1, 1]
Av = np.array([[2.0, 0, 0], [1, 3, 0], [1, 1, 4]], dtype="float32")
bv = np.array([1.0, 2.0, 3.0], dtype="float32")

print(pytensor.function([A, b], out, mode="CVM")(Av, bv))  # [1. 1. 1.]
print(pytensor.function([A, b], out, mode="MLX")(Av, bv))  # [0.5 0.5 0.5]

Potential fix (requires testing):

def solve_triangular(A, b):
    if op.unit_diagonal:
        eye = mx.eye(A.shape[-1], dtype=A.dtype)
        A = A * (1 - eye) + eye
    return mx.linalg.solve_triangular(..., upper=not lower, stream=mx.cpu)