Update JAX QR dispatch

Author: jessegrabowski

pytensor/link/jax/dispatch/nlinalg.py

@@ -9,7 +9,6 @@
     KroneckerProduct,
     MatrixInverse,
     MatrixPinv,
-    QRFull,
     SLogDet,
 )
 
@@ -67,16 +66,6 @@ def matrix_inverse(x):
     return matrix_inverse
 
 
-@jax_funcify.register(QRFull)
-def jax_funcify_QRFull(op, **kwargs):
-    mode = op.mode
-
-    def qr_full(x, mode=mode):
-        return jnp.linalg.qr(x, mode=mode)
-
-    return qr_full
-
-
 @jax_funcify.register(MatrixPinv)
 def jax_funcify_Pinv(op, **kwargs):
     def pinv(x):

pytensor/link/jax/dispatch/slinalg.py

@@ -5,6 +5,7 @@
 from pytensor.link.jax.dispatch.basic import jax_funcify
 from pytensor.tensor.slinalg import (
     LU,
+    QR,
     BlockDiagonal,
     Cholesky,
     CholeskySolve,
@@ -168,3 +169,13 @@ def cho_solve(c, b):
         )
 
     return cho_solve
+
+
+@jax_funcify.register(QR)
+def jax_funcify_QR(op, **kwargs):
+    mode = op.mode
+
+    def qr(x, mode=mode):
+        return jax.scipy.linalg.qr(x, mode=mode)
+
+    return qr

tests/link/jax/test_nlinalg.py

@@ -29,12 +29,6 @@ def assert_fn(x, y):
     outs = pt_nlinalg.eigh(x)
     compare_jax_and_py([x], outs, [X.astype(config.floatX)], assert_fn=assert_fn)
 
-    outs = pt_nlinalg.qr(x, mode="full")
-    compare_jax_and_py([x], outs, [X.astype(config.floatX)], assert_fn=assert_fn)
-
-    outs = pt_nlinalg.qr(x, mode="reduced")
-    compare_jax_and_py([x], outs, [X.astype(config.floatX)], assert_fn=assert_fn)
-
     outs = pt_nlinalg.svd(x)
     compare_jax_and_py([x], outs, [X.astype(config.floatX)], assert_fn=assert_fn)
 

tests/link/jax/test_slinalg.py

@@ -103,6 +103,18 @@ def test_jax_basic():
         ],
     )
 
+    def assert_fn(x, y):
+        np.testing.assert_allclose(x.astype(config.floatX), y, rtol=1e-3)
+
+    M = rng.normal(size=(3, 3))
+    X = M.dot(M.T)
+
+    outs = pt_slinalg.qr(x, mode="full")
+    compare_jax_and_py([x], outs, [X.astype(config.floatX)], assert_fn=assert_fn)
+
+    outs = pt_slinalg.qr(x, mode="economic")
+    compare_jax_and_py([x], outs, [X.astype(config.floatX)], assert_fn=assert_fn)
+
 
 def test_jax_solve():
     rng = np.random.default_rng(utt.fetch_seed())