Add slogdet for JAX

Author: mtsokol

pytensor/link/jax/dispatch/nlinalg.py

@@ -3,7 +3,7 @@
 from pytensor.link.jax.dispatch import jax_funcify
 from pytensor.tensor.blas import BatchedDot
 from pytensor.tensor.math import Dot, MaxAndArgmax
-from pytensor.tensor.nlinalg import SVD, Det, Eig, Eigh, MatrixInverse, QRFull
+from pytensor.tensor.nlinalg import SVD, Det, Eig, Eigh, MatrixInverse, QRFull, SLogDet
 
 
 @jax_funcify.register(SVD)
@@ -25,6 +25,14 @@ def det(x):
     return det
 
 
+@jax_funcify.register(SLogDet)
+def jax_funcify_SLogDet(op, **kwargs):
+    def slogdet(x):
+        return jnp.linalg.slogdet(x)
+
+    return slogdet
+
+
 @jax_funcify.register(Eig)
 def jax_funcify_Eig(op, **kwargs):
     def eig(x):

tests/link/jax/test_nlinalg.py

@@ -85,6 +85,10 @@ def assert_fn(x, y):
     out_fg = FunctionGraph([x], outs)
     compare_jax_and_py(out_fg, [X.astype(config.floatX)], assert_fn=assert_fn)
 
+    outs = at_nlinalg.slogdet(x)
+    out_fg = FunctionGraph([x], outs)
+    compare_jax_and_py(out_fg, [X.astype(config.floatX)], assert_fn=assert_fn)
+
 
 @pytest.mark.xfail(
     version_parse(jax.__version__) >= version_parse("0.2.12"),