Add jax dispatch for lu_solve

Author: jessegrabowski

pytensor/link/jax/dispatch/slinalg.py

@@ -9,6 +9,7 @@
     Cholesky,
     Eigvalsh,
     LUFactor,
+    LUSolve,
     Solve,
     SolveTriangular,
 )
@@ -119,9 +120,27 @@ def jax_funcify_LUFactor(op, **kwargs):
     check_finite = op.check_finite
     overwrite_a = op.overwrite_a
 
-    def lu_factor(*inputs):
+    def lu_factor(a):
         return jax.scipy.linalg.lu_factor(
-            *inputs, check_finite=check_finite, overwrite_a=overwrite_a
+            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

@@ -12,7 +12,6 @@
 from pytensor.tensor import subtensor as pt_subtensor
 from pytensor.tensor.math import clip, cosh
 from pytensor.tensor.type import matrix, vector
-from tests import unittest_tools as utt
 from tests.link.jax.test_basic import compare_jax_and_py
 
 
@@ -272,3 +271,17 @@ def test_jax_lu_factor(shape):
         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])