Add JAX dispatch for LU

Author: jessegrabowski

pytensor/link/jax/dispatch/slinalg.py

@@ -2,6 +2,7 @@
 
 from pytensor.link.jax.dispatch.basic import jax_funcify
 from pytensor.tensor.slinalg import (
+    LU,
     BlockDiagonal,
     Cholesky,
     Eigvalsh,
@@ -76,3 +77,20 @@ def block_diag(*inputs):
         return jax.scipy.linalg.block_diag(*inputs)
 
     return block_diag
+
+
+@jax_funcify.register(LU)
+def jax_funcify_LU(op, **kwargs):
+    permute_l = op.permute_l
+    p_indices = op.p_indices
+    check_finite = op.check_finite
+
+    if p_indices:
+        raise ValueError("JAX does not support the p_indices argument")
+
+    def lu(*inputs):
+        return jax.scipy.linalg.lu(
+            *inputs, permute_l=permute_l, check_finite=check_finite
+        )
+
+    return lu

tests/link/jax/test_slinalg.py

@@ -219,3 +219,33 @@ def test_jax_solve_discrete_lyapunov(
         jax_mode="JAX",
         assert_fn=partial(np.testing.assert_allclose, atol=atol, rtol=rtol),
     )
+
+
+@pytest.mark.parametrize(
+    "permute_l, p_indices",
+    [(True, False), (False, True), (False, False)],
+    ids=["PL", "p_indices", "P"],
+)
+@pytest.mark.parametrize("complex", [False, True], ids=["real", "complex"])
+@pytest.mark.parametrize("shape", [(3, 5, 5), (5, 5)], ids=["batched", "not_batched"])
+def test_jax_lu(permute_l, p_indices, complex, shape: tuple[int]):
+    rng = np.random.default_rng()
+    A = pt.tensor(
+        "A",
+        shape=shape,
+        dtype=f"complex{int(config.floatX[-2:]) * 2}" if complex else config.floatX,
+    )
+    out = pt_slinalg.lu(A, permute_l=permute_l, p_indices=p_indices)
+    out_fg = FunctionGraph([A], out)
+
+    x = rng.normal(size=shape).astype(config.floatX)
+    if complex:
+        x = x + 1j * rng.normal(size=shape).astype(config.floatX)
+
+    if p_indices:
+        with pytest.raises(
+            ValueError, match="JAX does not support the p_indices argument"
+        ):
+            compare_jax_and_py(out_fg, [x])
+    else:
+        compare_jax_and_py(out_fg, [x])