Add JAX dispatch for expm

Author: Jesse Grabowski

pytensor/link/jax/dispatch/slinalg.py

@@ -10,6 +10,7 @@
     Cholesky,
     CholeskySolve,
     Eigvalsh,
+    Expm,
     LUFactor,
     PivotToPermutations,
     Solve,
@@ -179,3 +180,11 @@ def qr(x, mode=mode):
         return jax.scipy.linalg.qr(x, mode=mode)
 
     return qr
+
+
+@jax_funcify.register(Expm)
+def jax_funcify_Expm(op, **kwargs):
+    def expm(x):
+        return jax.scipy.linalg.expm(x)
+
+    return expm

tests/link/jax/test_slinalg.py

@@ -361,3 +361,12 @@ def test_jax_cho_solve(b_shape, lower):
     out = pt_slinalg.cho_solve((c, lower), b, b_ndim=len(b_shape))
 
     compare_jax_and_py([A, b], [out], [A_val, b_val])
+
+
+def test_jax_expm():
+    rng = np.random.default_rng(utt.fetch_seed())
+    A = pt.tensor(name="A", shape=(5, 5))
+    A_val = rng.normal(size=(5, 5)).astype(config.floatX)
+    out = pt_slinalg.expm(A)
+
+    compare_jax_and_py([A], [out], [A_val])