JAX dispatches for LU Ops

Author: jessegrabowski

pytensor/link/jax/dispatch/slinalg.py

@@ -4,9 +4,12 @@
 
 from pytensor.link.jax.dispatch.basic import jax_funcify
 from pytensor.tensor.slinalg import (
+    LU,
     BlockDiagonal,
     Cholesky,
     Eigvalsh,
+    LUFactor,
+    LUSolve,
     Solve,
     SolveTriangular,
 )
@@ -93,3 +96,51 @@ 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
+
+
+@jax_funcify.register(LUFactor)
+def jax_funcify_LUFactor(op, **kwargs):
+    check_finite = op.check_finite
+    overwrite_a = op.overwrite_a
+
+    def lu_factor(a):
+        return jax.scipy.linalg.lu_factor(
+            a, check_finite=check_finite, overwrite_a=overwrite_a
+        )
+
+    return lu_factor
+
+
+@jax_funcify.register(LUSolve)
+def jax_funcify_LUSolve(op, **kwargs):
+    trans = op.trans
+    check_finite = op.check_finite
+    overwrite_b = op.overwrite_b
+
+    def lu_solve(lu, pivots, b):
+        return jax.scipy.linalg.lu_solve(
+            (lu, pivots),
+            b,
+            trans=trans,
+            check_finite=check_finite,
+            overwrite_b=overwrite_b,
+        )
+
+    return lu_solve

tests/link/jax/test_slinalg.py

@@ -228,3 +228,60 @@ 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)
+
+    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(graph_inputs=[A], graph_outputs=out, test_inputs=[x])
+    else:
+        compare_jax_and_py(graph_inputs=[A], graph_outputs=out, test_inputs=[x])
+
+
+@pytest.mark.parametrize("shape", [(5, 5), (5, 5, 5)], ids=["matrix", "batch"])
+def test_jax_lu_factor(shape):
+    rng = np.random.default_rng(utt.fetch_seed())
+    A = pt.tensor(name="A", shape=shape)
+    A_value = rng.normal(size=shape).astype(config.floatX)
+    out = pt_slinalg.lu_factor(A)
+
+    compare_jax_and_py(
+        [A],
+        out,
+        [A_value],
+    )
+
+
+@pytest.mark.parametrize("b_shape", [(5,), (5, 5)])
+def test_jax_lu_solve(b_shape):
+    rng = np.random.default_rng(utt.fetch_seed())
+    A_val = rng.normal(size=(5, 5)).astype(config.floatX)
+    b_val = rng.normal(size=b_shape).astype(config.floatX)
+
+    A = pt.tensor(name="A", shape=(5, 5))
+    b = pt.tensor(name="b", shape=b_shape)
+    lu_and_pivots = pt_slinalg.lu_factor(A)
+    out = pt_slinalg.lu_solve(lu_and_pivots, b)
+
+    compare_jax_and_py([A, b], [out], [A_val, b_val])