Remove unnecessary aesara.tensor.nlinalg.QRIncomplete

aesara/gpuarray/opt.py

@@ -2887,23 +2887,6 @@ def local_gpu_magma_qr(fgraph, op, context_name, inputs, outputs):
     return out
 
 
-@register_opt("magma", "fast_compile")
-@op_lifter([nlinalg.QRIncomplete])
-@register_opt2([nlinalg.QRIncomplete], "magma", "fast_compile")
-def local_gpu_magma_qr_incomplete(fgraph, op, context_name, inputs, outputs):
-    if not config.magma__enabled:
-        return
-    if inputs[0].dtype not in ["float16", "float32"]:
-        return
-    x = inputs[0]
-    if inputs[0].dtype == "float16":
-        x = inputs[0].astype("float32")
-    out = gpu_qr(x, complete=False)
-    if inputs[0].dtype == "float16":
-        return [out.astype("float16")]
-    return out
-
-
 # Matrix inverse
 @register_opt("magma", "fast_compile")
 @op_lifter([nlinalg.MatrixInverse])

aesara/link/jax/dispatch.py

@@ -44,15 +44,7 @@ from aesara.tensor.extra_ops import (
     UnravelIndex,
 )
 from aesara.tensor.math import Dot, MaxAndArgmax
-from aesara.tensor.nlinalg import (
-    SVD,
-    Det,
-    Eig,
-    Eigh,
-    MatrixInverse,
-    QRFull,
-    QRIncomplete,
-)
+from aesara.tensor.nlinalg import SVD, Det, Eig, Eigh, MatrixInverse, QRFull
 from aesara.tensor.nnet.basic import LogSoftmax, Softmax
 from aesara.tensor.random.op import RandomVariable
 from aesara.tensor.shape import Reshape, Shape, Shape_i, SpecifyShape
@@ -835,16 +827,6 @@ def jax_funcify_QRFull(op, **kwargs):
     return qr_full
 
 
-@jax_funcify.register(QRIncomplete)
-def jax_funcify_QRIncomplete(op, **kwargs):
-    mode = op.mode
-
-    def qr_incomplete(x, mode=mode):
-        return jnp.linalg.qr(x, mode=mode)
-
-    return qr_incomplete
-
-
 @jax_funcify.register(SVD)
 def jax_funcify_SVD(op, **kwargs):
     full_matrices = op.full_matrices

aesara/tensor/nlinalg.py

@@ -408,47 +408,31 @@ class QRFull(Op):
     def make_node(self, x):
         x = as_tensor_variable(x)
         assert x.ndim == 2, "The input of qr function should be a matrix."
-        q = matrix(dtype=x.dtype)
+
+        in_dtype = x.type.numpy_dtype
+        out_dtype = np.dtype(f"f{in_dtype.itemsize}")
+
         if self.mode != "raw":
             r = matrix(dtype=x.dtype)
         else:
             r = vector(dtype=x.dtype)
 
-        return Apply(self, [x], [q, r])
-
-    def perform(self, node, inputs, outputs):
-        (x,) = inputs
-        (q, r) = outputs
-        assert x.ndim == 2, "The input of qr function should be a matrix."
-        q[0], r[0] = self._numop(x, self.mode)
-
-
-class QRIncomplete(Op):
-    """
-    Incomplete QR Decomposition.
-
-    Computes the QR decomposition of a matrix.
-    Factor the matrix a as qr and return a single matrix R.
-
-    """
-
-    _numop = staticmethod(np.linalg.qr)
-    __props__ = ("mode",)
-
-    def __init__(self, mode):
-        self.mode = mode
+        if self.mode != "r":
+            q = matrix(dtype=out_dtype)
+            outputs = [q, r]
+        else:
+            outputs = [r]
 
-    def make_node(self, x):
-        x = as_tensor_variable(x)
-        assert x.ndim == 2, "The input of qr function should be a matrix."
-        r = matrix(dtype=x.dtype)
-        return Apply(self, [x], [r])
+        return Apply(self, [x], outputs)
 
     def perform(self, node, inputs, outputs):
         (x,) = inputs
-        (r,) = outputs
         assert x.ndim == 2, "The input of qr function should be a matrix."
-        r[0] = self._numop(x, self.mode)
+        res = self._numop(x, self.mode)
+        if self.mode != "r":
+            outputs[0][0], outputs[1][0] = res
+        else:
+            outputs[0][0] = res
 
 
 def qr(a, mode="reduced"):
@@ -493,12 +477,7 @@ def qr(a, mode="reduced"):
         The upper-triangular matrix.
 
     """
-
-    x = [[2, 1], [3, 4]]
-    if isinstance(np.linalg.qr(x, mode), tuple):
-        return QRFull(mode)(a)
-    else:
-        return QRIncomplete(mode)(a)
+    return QRFull(mode)(a)
 
 
 class SVD(Op):

tests/gpuarray/test_linalg.py

@@ -22,15 +22,7 @@ from aesara.gpuarray.linalg import (
     gpu_solve_lower_triangular,
     gpu_svd,
 )
-from aesara.tensor.nlinalg import (
-    SVD,
-    MatrixInverse,
-    QRFull,
-    QRIncomplete,
-    eigh,
-    matrix_inverse,
-    qr,
-)
+from aesara.tensor.nlinalg import SVD, MatrixInverse, QRFull, eigh, matrix_inverse, qr
 from aesara.tensor.slinalg import Cholesky, cholesky, imported_scipy
 from aesara.tensor.type import fmatrix, matrix, tensor3, vector
 from tests import unittest_tools as utt
@@ -380,7 +372,6 @@ class TestMagma:
             (MatrixInverse(), GpuMagmaMatrixInverse),
             (SVD(), GpuMagmaSVD),
             (QRFull(mode="reduced"), GpuMagmaQR),
-            (QRIncomplete(mode="r"), GpuMagmaQR),
             # TODO: add support for float16 to Eigh numpy
             # (Eigh(), GpuMagmaEigh),
             (Cholesky(), GpuMagmaCholesky),