Update Torch QR dispatch

Author: jessegrabowski

pytensor/link/pytorch/dispatch/__init__.py

@@ -8,6 +8,7 @@
 import pytensor.link.pytorch.dispatch.math
 import pytensor.link.pytorch.dispatch.extra_ops
 import pytensor.link.pytorch.dispatch.nlinalg
+import pytensor.link.pytorch.dispatch.slinalg
 import pytensor.link.pytorch.dispatch.shape
 import pytensor.link.pytorch.dispatch.sort
 import pytensor.link.pytorch.dispatch.subtensor

pytensor/link/pytorch/dispatch/nlinalg.py

@@ -9,7 +9,6 @@
     KroneckerProduct,
     MatrixInverse,
     MatrixPinv,
-    QRFull,
     SLogDet,
 )
 
@@ -70,21 +69,6 @@ def matrix_inverse(x):
     return matrix_inverse
 
 
-@pytorch_funcify.register(QRFull)
-def pytorch_funcify_QRFull(op, **kwargs):
-    mode = op.mode
-    if mode == "raw":
-        raise NotImplementedError("raw mode not implemented in PyTorch")
-
-    def qr_full(x):
-        Q, R = torch.linalg.qr(x, mode=mode)
-        if mode == "r":
-            return R
-        return Q, R
-
-    return qr_full
-
-
 @pytorch_funcify.register(MatrixPinv)
 def pytorch_funcify_Pinv(op, **kwargs):
     hermitian = op.hermitian

pytensor/link/pytorch/dispatch/slinalg.py

@@ -0,0 +1,19 @@
+import torch
+
+from pytensor.link.pytorch.dispatch import pytorch_funcify
+from pytensor.tensor.slinalg import QR
+
+
+@pytorch_funcify.register(QR)
+def pytorch_funcify_QR(op, **kwargs):
+    mode = op.mode
+    if mode == "raw":
+        raise NotImplementedError("raw mode not implemented in PyTorch")
+
+    def qr(x):
+        Q, R = torch.linalg.qr(x, mode=mode)
+        if mode == "r":
+            return R
+        return Q, R
+
+    return qr

tests/link/pytorch/conftest.py

@@ -0,0 +1,16 @@
+import numpy as np
+import pytest
+
+from pytensor import config
+from pytensor.tensor.type import matrix
+
+
+@pytest.fixture
+def matrix_test():
+    rng = np.random.default_rng(213234)
+
+    M = rng.normal(size=(3, 3))
+    test_value = M.dot(M.T).astype(config.floatX)
+
+    x = matrix("x")
+    return x, test_value

tests/link/pytorch/test_nlinalg.py

@@ -8,17 +8,6 @@
 from tests.link.pytorch.test_basic import compare_pytorch_and_py
 
 
-@pytest.fixture
-def matrix_test():
-    rng = np.random.default_rng(213234)
-
-    M = rng.normal(size=(3, 3))
-    test_value = M.dot(M.T).astype(config.floatX)
-
-    x = matrix("x")
-    return (x, test_value)
-
-
 @pytest.mark.parametrize(
     "func",
     (pt_nla.eig, pt_nla.eigh, pt_nla.SLogDet(), pt_nla.inv, pt_nla.det),
@@ -34,22 +23,6 @@ def assert_fn(x, y):
     compare_pytorch_and_py([x], outs, [test_value], assert_fn=assert_fn)
 
 
-@pytest.mark.parametrize(
-    "mode",
-    (
-        "complete",
-        "reduced",
-        "r",
-        pytest.param("raw", marks=pytest.mark.xfail(raises=NotImplementedError)),
-    ),
-)
-def test_qr(mode, matrix_test):
-    x, test_value = matrix_test
-    outs = pt_nla.qr(x, mode=mode)
-
-    compare_pytorch_and_py([x], outs, [test_value])
-
-
 @pytest.mark.parametrize("compute_uv", [True, False])
 @pytest.mark.parametrize("full_matrices", [True, False])
 def test_svd(compute_uv, full_matrices, matrix_test):

tests/link/pytorch/test_slinalg.py

@@ -0,0 +1,20 @@
+import pytest
+
+import pytensor
+from tests.link.pytorch.test_basic import compare_pytorch_and_py
+
+
+@pytest.mark.parametrize(
+    "mode",
+    (
+        "complete",
+        "reduced",
+        "r",
+        pytest.param("raw", marks=pytest.mark.xfail(raises=NotImplementedError)),
+    ),
+)
+def test_qr(mode, matrix_test):
+    x, test_value = matrix_test
+    outs = pytensor.tensor.slinalg.qr(x, mode=mode)
+
+    compare_pytorch_and_py([x], outs, [test_value])