Move QR helper functions to qr_helper.py

vlad-perevezentsev · vlad-perevezentsev · commit 707c5a0fa39c · 2026-03-09T05:32:50.000-07:00
diff --git a/dpnp/tests/qr_helper.py b/dpnp/tests/qr_helper.py
@@ -0,0 +1,70 @@
+import numpy
+
+from .helper import has_support_aspect64
+
+
+def gram(x, xp):
+    # Return Gram matrix: X^H @ X
+    return xp.conjugate(x).swapaxes(-1, -2) @ x
+
+
+def get_R_from_raw(h, m, n, xp):
+    # Get reduced R from NumPy-style raw QR:
+    # R = triu((tril(h))^T), shape (..., k, n)
+    k = min(m, n)
+    rt = xp.tril(h)
+    r = xp.swapaxes(rt, -1, -2)
+    r = xp.triu(r[..., :m, :n])
+    return r[..., :k, :]
+
+
+def check_qr(a_np, a_xp, mode, xp):
+    # QR is not unique:
+    # element-wise comparison with NumPy may differ by sign/phase.
+    # To verify correctness use mode-dependent functional checks:
+    # complete/reduced: check decomposition Q @ R = A
+    # raw/r: check invariant R^H @ R = A^H @ A
+    if mode in ("complete", "reduced"):
+        res = xp.linalg.qr(a_xp, mode)
+        assert xp.allclose(res.Q @ res.R, a_xp, atol=1e-5)
+
+    # Since QR satisfies A = Q @ R with orthonormal Q (Q^H @ Q = I),
+    # validate correctness via the invariant R^H @ R == A^H @ A
+    # for raw/r modes
+    elif mode == "raw":
+        _, tau_np = numpy.linalg.qr(a_np, mode=mode)
+        h_xp, tau_xp = xp.linalg.qr(a_xp, mode=mode)
+
+        m, n = a_np.shape[-2], a_np.shape[-1]
+        Rraw_xp = get_R_from_raw(h_xp, m, n, xp)
+
+        # Use reduced QR as a reference:
+        # reduced is validated via Q @ R == A
+        exp_res = xp.linalg.qr(a_xp, mode="reduced")
+        exp_r = exp_res.R
+        assert xp.allclose(Rraw_xp, exp_r, atol=1e-4, rtol=1e-4)
+
+        exp_xp = gram(a_xp, xp)
+
+        # Compare R^H @ R == A^H @ A
+        assert xp.allclose(gram(Rraw_xp, xp), exp_xp, atol=1e-4, rtol=1e-4)
+
+        assert tau_xp.shape == tau_np.shape
+        if not has_support_aspect64(tau_xp.sycl_device):
+            assert tau_xp.dtype.kind == tau_np.dtype.kind
+        else:
+            assert tau_xp.dtype == tau_np.dtype
+
+    else:  # mode == "r"
+        r_xp = xp.linalg.qr(a_xp, mode="r")
+
+        # Use reduced QR as a reference:
+        # reduced is validated via Q @ R == A
+        exp_res = xp.linalg.qr(a_xp, mode="reduced")
+        exp_r = exp_res.R
+        assert xp.allclose(r_xp, exp_r, atol=1e-4, rtol=1e-4)
+
+        exp_xp = gram(a_xp, xp)
+
+        # Compare R^H @ R == A^H @ A
+        assert xp.allclose(gram(r_xp, xp), exp_xp, atol=1e-4, rtol=1e-4)
diff --git a/dpnp/tests/test_linalg.py b/dpnp/tests/test_linalg.py
@@ -25,6 +25,7 @@
     has_support_aspect64,
     numpy_version,
 )
+from .qr_helper import check_qr
 from .third_party.cupy import testing
 
 
@@ -3135,85 +3136,6 @@ def test_error(self):
 
 
 class TestQr:
-    def gram(self, X, xp):
-        # Return Gram matrix: X^H @ X
-        return xp.conjugate(X).swapaxes(-1, -2) @ X
-
-    def get_R_from_raw(self, h, m, n, xp):
-        # Get reduced R from NumPy-style raw QR:
-        # R = triu((tril(h))^T), shape (..., k, n)
-        k = min(m, n)
-        Rt = xp.tril(h)
-        R = xp.swapaxes(Rt, -1, -2)
-        R = xp.triu(R[..., :m, :n])
-
-        return R[..., :k, :]
-
-    # QR is not unique:
-    # element-wise comparison with NumPy may differ by sign/phase.
-    # To verify correctness use mode-dependent functional checks:
-    # complete/reduced: check decomposition Q @ R = A
-    # raw/r: check invariant R^H @ R = A^H @ A
-    def check_qr(self, a_np, a_dp, mode):
-        if mode in ("complete", "reduced"):
-            res = dpnp.linalg.qr(a_dp, mode)
-            assert dpnp.allclose(res.Q @ res.R, a_dp, atol=1e-5)
-
-        # Since QR satisfies A = Q @ R with orthonormal Q (Q^H @ Q = I),
-        # validate correctness via the invariant R^H @ R == A^H @ A
-        # for raw/r modes
-        elif mode == "raw":
-            h_np, tau_np = numpy.linalg.qr(a_np, mode=mode)
-            h_dp, tau_dp = dpnp.linalg.qr(a_dp, mode=mode)
-
-            m, n = a_np.shape[-2], a_np.shape[-1]
-            Rraw_np = self.get_R_from_raw(h_np, m, n, numpy)
-            Rraw_dp = self.get_R_from_raw(h_dp, m, n, dpnp)
-
-            # Use reduced QR as a reference:
-            # reduced is validated via Q @ R == A
-            exp_res = dpnp.linalg.qr(a_dp, mode="reduced")
-            exp_R = exp_res.R
-            assert_allclose(Rraw_dp, exp_R, atol=1e-4, rtol=1e-4)
-
-            exp_dp = self.gram(a_dp, dpnp).astype(Rraw_dp.dtype)
-            exp_np = self.gram(a_np, numpy).astype(Rraw_np.dtype)
-
-            # compare R^H @ R == A^H @ A
-            assert_allclose(
-                self.gram(Rraw_dp, dpnp), exp_dp, atol=1e-4, rtol=1e-4
-            )
-            assert_allclose(
-                self.gram(Rraw_np, numpy), exp_np, atol=1e-4, rtol=1e-4
-            )
-
-            assert tau_dp.shape == tau_np.shape
-            if not has_support_aspect64(tau_dp.sycl_device):
-                if tau_np.dtype == numpy.float64:
-                    tau_np = tau_np.astype("float32")
-                elif tau_np.dtype == numpy.complex128:
-                    tau_np = tau_np.astype("complex64")
-            assert tau_dp.dtype == tau_np.dtype
-
-        else:  # mode == "r"
-            R_np = numpy.linalg.qr(a_np, mode="r")
-            R_dp = dpnp.linalg.qr(a_dp, mode="r")
-
-            # Use reduced QR as a reference:
-            # reduced is validated via Q @ R == A
-            exp_res = dpnp.linalg.qr(a_dp, mode="reduced")
-            exp_R = exp_res.R
-            assert_allclose(R_dp, exp_R, atol=1e-4, rtol=1e-4)
-
-            exp_dp = self.gram(a_dp, dpnp)
-            exp_np = self.gram(a_np, numpy)
-
-            # compare R^H @ R == A^H @ A
-            assert_allclose(self.gram(R_dp, dpnp), exp_dp, atol=1e-4, rtol=1e-4)
-            assert_allclose(
-                self.gram(R_np, numpy), exp_np, atol=1e-4, rtol=1e-4
-            )
-
     @pytest.mark.parametrize("dtype", get_float_complex_dtypes())
     @pytest.mark.parametrize(
         "shape",
@@ -3245,7 +3167,7 @@ def test_qr(self, dtype, shape, mode):
         a = generate_random_numpy_array(shape, dtype, seed_value=None)
         ia = dpnp.array(a, dtype=dtype)
 
-        self.check_qr(a, ia, mode)
+        check_qr(a, ia, mode, dpnp)
 
     @pytest.mark.parametrize("dtype", get_float_complex_dtypes())
     @pytest.mark.parametrize(
@@ -3258,7 +3180,7 @@ def test_qr_large(self, dtype, shape, mode):
         a = generate_random_numpy_array(shape, dtype, seed_value=81)
         ia = dpnp.array(a)
 
-        self.check_qr(a, ia, mode)
+        check_qr(a, ia, mode, dpnp)
 
     @pytest.mark.parametrize("dtype", get_float_complex_dtypes())
     @pytest.mark.parametrize(
@@ -3278,17 +3200,17 @@ def test_qr_empty(self, dtype, shape, mode):
         a = numpy.empty(shape, dtype=dtype)
         ia = dpnp.array(a)
 
-        self.check_qr(a, ia, mode)
+        check_qr(a, ia, mode, dpnp)
 
     @pytest.mark.parametrize("mode", ["complete", "reduced", "r", "raw"])
     def test_qr_strides(self, mode):
         a = generate_random_numpy_array((5, 5))
         ia = dpnp.array(a)
 
         # positive strides
-        self.check_qr(a[::2, ::2], ia[::2, ::2], mode)
+        check_qr(a[::2, ::2], ia[::2, ::2], mode, dpnp)
         # negative strides
-        self.check_qr(a[::-2, ::-2], ia[::-2, ::-2], mode)
+        check_qr(a[::-2, ::-2], ia[::-2, ::-2], mode, dpnp)
 
     def test_qr_errors(self):
         a_dp = dpnp.array([[1, 2], [3, 5]], dtype="float32")
diff --git a/dpnp/tests/third_party/cupy/linalg_tests/test_decomposition.py b/dpnp/tests/third_party/cupy/linalg_tests/test_decomposition.py
@@ -14,6 +14,7 @@
 from dpnp.tests.helper import (
     has_support_aspect64,
 )
+from dpnp.tests.qr_helper import check_qr
 from dpnp.tests.third_party.cupy import testing
 from dpnp.tests.third_party.cupy.testing import _condition
 
@@ -167,19 +168,6 @@ def test_decomposition(self, dtype):
     )
 )
 class TestQRDecomposition(unittest.TestCase):
-    def _gram(self, x, xp):
-        # Gram matrix: X^H @ X
-        return xp.conjugate(x).swapaxes(-1, -2) @ x
-
-    def _get_R_from_raw(self, h, m, n, xp):
-        # Get reduced R from NumPy-style raw QR:
-        # R = triu((tril(h))^T), shape (..., k, n)
-        k = min(m, n)
-        Rt = xp.tril(h)
-        R = xp.swapaxes(Rt, -1, -2)
-        R = xp.triu(R[..., :m, :n])
-        return R[..., :k, :]
-
     @testing.for_dtypes("fdFD")
     def check_mode(self, array, mode, dtype):
         # if runtime.is_hip and driver.get_build_version() < 307:
@@ -188,13 +176,20 @@ def check_mode(self, array, mode, dtype):
 
         a_cpu = numpy.asarray(array, dtype=dtype)
         a_gpu = cupy.asarray(array, dtype=dtype)
-        result_gpu = cupy.linalg.qr(a_gpu, mode=mode)
+        # QR is not unique:
+        # element-wise comparison with NumPy may differ by sign/phase.
+        # To verify correctness use mode-dependent functional checks:
+        # complete/reduced: check decomposition Q @ R = A
+        # raw/r: check invariant R^H @ R = A^H @ A
+
+        # result_gpu = cupy.linalg.qr(a_gpu, mode=mode)
         if (
             mode != "raw"
             or numpy.lib.NumpyVersion(numpy.__version__) >= "1.22.0rc1"
         ):
-            result_cpu = numpy.linalg.qr(a_cpu, mode=mode)
-            self._check_result(result_cpu, result_gpu, a_cpu, a_gpu, mode)
+            # result_cpu = numpy.linalg.qr(a_cpu, mode=mode)
+            # self._check_result(result_cpu, result_gpu, a_gpu, mode)
+            check_qr(a_cpu, a_gpu, mode, cupy)
 
     # def _check_result(self, result_cpu, result_gpu):
     #     if isinstance(result_cpu, tuple):
@@ -205,54 +200,6 @@ def check_mode(self, array, mode, dtype):
     #         assert result_cpu.dtype == result_gpu.dtype
     #         testing.assert_allclose(result_cpu, result_gpu, atol=1e-4)
 
-    # QR is not unique:
-    # element-wise comparison with NumPy may differ by sign/phase.
-    # To verify correctness use mode-dependent functional checks:
-    # complete/reduced: check decomposition Q @ R = A
-    # raw/r: check invariant R^H @ R = A^H @ A
-    def _check_result(self, result_cpu, result_gpu, a_cpu, a_gpu, mode):
-
-        if mode in ("complete", "reduced"):
-            q_gpu, r_gpu = result_gpu
-            testing.assert_allclose(q_gpu @ r_gpu, a_gpu, atol=1e-4)
-
-        elif mode == "raw":
-            h_gpu, tau_gpu = result_gpu
-            h_cpu, tau_cpu = result_cpu
-            m, n = a_gpu.shape[-2], a_gpu.shape[-1]
-            r_gpu = self._get_R_from_raw(h_gpu, m, n, cupy)
-            r_cpu = self._get_R_from_raw(h_cpu, m, n, numpy)
-
-            exp_gpu = self._gram(a_gpu, cupy)
-            exp_cpu = self._gram(a_cpu, numpy)
-
-            testing.assert_allclose(
-                self._gram(r_gpu, cupy), exp_gpu, atol=1e-4, rtol=1e-4
-            )
-            testing.assert_allclose(
-                self._gram(r_cpu, numpy), exp_cpu, atol=1e-4, rtol=1e-4
-            )
-
-            assert tau_gpu.shape == tau_cpu.shape
-            if not has_support_aspect64(tau_gpu.sycl_device):
-                if tau_cpu.dtype == numpy.float64:
-                    tau_cpu = tau_cpu.astype("float32")
-                elif tau_cpu.dtype == numpy.complex128:
-                    tau_cpu = tau_cpu.astype("complex64")
-            assert tau_gpu.dtype == tau_cpu.dtype
-
-        else:  # mode == "r"
-            r_gpu = result_gpu
-            r_cpu = result_cpu
-            exp_gpu = self._gram(a_gpu, cupy)
-            exp_cpu = self._gram(a_cpu, numpy)
-            testing.assert_allclose(
-                self._gram(r_gpu, cupy), exp_gpu, atol=1e-4, rtol=1e-4
-            )
-            testing.assert_allclose(
-                self._gram(r_cpu, numpy), exp_cpu, atol=1e-4, rtol=1e-4
-            )
-
     @testing.fix_random()
     @_condition.repeat(3, 10)
     def test_mode(self):
