MLX: batched solve fails with "[Primitive::vmap] Not implemented for LUF"

[jessegrabowski]

Blockwise lowers to mx.vmap, which has no rule for the LUF primitive, so every batched solve fails. MLX's own CPU-stream mx.linalg.solve and lu_factor already accept batched input, so the vmap wrapper isn't needed for these in the first place.

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

A = pt.tensor("A", shape=(4, 3, 3), dtype="float32")
b = pt.tensor("b", shape=(4, 3), dtype="float32")
Av = np.broadcast_to(np.eye(3, dtype="float32") * 2, (4, 3, 3)).copy()
bv = np.ones((4, 3), dtype="float32")

out = pt.linalg.solve(A, b, b_ndim=1)
print(pytensor.function([A, b], out, mode="CVM")(Av, bv)[0])  # [0.5 0.5 0.5]
print(pytensor.function([A, b], out, mode="MLX")(Av, bv)[0])  # ValueError: [Primitive::vmap] Not implemented for LUF.

Batched cholesky, solve_triangular and inv all work; solve and lu_factor are the ones that hit this.