Repository: cupy/cupy

Issue:
Batched QR solver
This is for ~~both NumPy compatibility and also~~ the Python Array API standard (#4789, WIP proposal: data-apis/array-api#126).

Currently we have this check:
https://github.com/cupy/cupy/blob/07171d785422df01d2d0b41b41268c26d02d1331/cupy/linalg/_decomposition.py#L238-L240
so we can only support a 2D matrix, but we might be able to follow the approach of batched svd and pinv to support batched QR factorization. I can look into this later.


Discussion:
Oops, read it too quickly: NumPy does not support it currently, so this is just for Array API.
```python
>>> import numpy as np
>>> a = np.random.random((2,16,16))
>>> b = np.linalg.qr(a)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<__array_function__ internals>", line 6, in qr
  File "/home/leofang/miniconda3/envs/cupy_cuda112_dev/lib/python3.7/site-packages/numpy/linalg/linalg.py", line 907, in qr
    _assert_2d(a)
  File "/home/leofang/miniconda3/envs/cupy_cuda112_dev/lib/python3.7/site-packages/numpy/linalg/linalg.py", line 191, in _assert_2d
    'two-dimensional' % a.ndim)
numpy.linalg.LinAlgError: 3-dimensional array given. Array must be two-dimensional
```
Update: In the next NumPy release batched QR will be supported, see https://github.com/numpy/numpy/pull/19151. We should be prepared.


Resolving patch:
diff --git a/cupy/linalg/_decomposition.py b/cupy/linalg/_decomposition.py
--- a/cupy/linalg/_decomposition.py
+++ b/cupy/linalg/_decomposition.py
@@ -8,6 +8,7 @@
 from cupy.cuda import device
 from cupy.cusolver import check_availability
 from cupy.cusolver import _gesvdj_batched, _gesvd_batched
+from cupy.cusolver import _geqrf_orgqr_batched
 from cupy.linalg import _util
 
 
@@ -210,6 +211,48 @@ def cholesky(a):
     return x.astype(out_dtype, copy=False)
 
 
+def _qr_batched(a, mode):
+    batch_shape = a.shape[:-2]
+    batch_size = internal.prod(batch_shape)
+    m, n = a.shape[-2:]
+
+    # first handle any 0-size inputs
+    if batch_size == 0 or m == 0 or n == 0:
+        # support float32, float64, complex64, and complex128
+        dtype, out_dtype = _util.linalg_common_type(a)
+        if mode == 'raw':
+            # compatibility with numpy.linalg.qr
+            out_dtype = numpy.promote_types(out_dtype, 'd')
+
+        if mode == 'reduced':
+            return (cupy.empty(batch_shape + (m, 0), out_dtype),
+                    cupy.empty(batch_shape + (0, n), out_dtype))
+        elif mode == 'complete':
+            if batch_size == 0:
+                q = cupy.empty(batch_shape + (m, m), out_dtype)
+            else:
+                q = cupy.eye(m)
+                q = cupy.stack([q for i in range(batch_size)])
+            return (q, cupy.empty(batch_shape + (m, n), out_dtype))
+        elif mode == 'r':
+            return cupy.empty(batch_shape + (0, n), out_dtype)
+        elif mode == 'raw':
+            return (cupy.empty(batch_shape + (n, m), out_dtype),
+                    cupy.empty(batch_shape + (0,), out_dtype))
+
+    # ...then delegate real computation to cuSOLVER/rocSOLVER
+    a = a.reshape(-1, *(a.shape[-2:]))
+    out = _geqrf_orgqr_batched(a, mode)
+
+    if mode == 'r':
+        return out.reshape(batch_shape + out.shape[-2:])
+    q, r = out
+    q = q.reshape(batch_shape + q.shape[-2:])
+    idx = -1 if mode == 'raw' else -2
+    r.reshape(batch_shape + r.shape[idx:])
+    return (q, r)
+
+
 def qr(a, mode='reduced'):
     """QR decomposition.
 
@@ -220,9 +263,9 @@ def qr(a, mode='reduced'):
         a (cupy.ndarray): The input matrix.
         mode (str): The mode of decomposition. Currently 'reduced',
             'complete', 'r', and 'raw' modes are supported. The default mode
-            is 'reduced', in which matrix ``A = (M, N)`` is decomposed into
-            ``Q``, ``R`` with dimensions ``(M, K)``, ``(K, N)``, where
-            ``K = min(M, N)``.
+            is 'reduced', in which matrix ``A = (..., M, N)`` is decomposed
+            into ``Q``, ``R`` with dimensions ``(..., M, K)``, ``(..., K, N)``,
+            where ``K = min(M, N)``.
 
     Returns:
         cupy.ndarray, or tuple of ndarray:
@@ -239,16 +282,17 @@ def qr(a, mode='reduced'):
 
     .. seealso:: :func:`numpy.linalg.qr`
     """
-    # TODO(Saito): Current implementation only accepts two-dimensional arrays
     _util._assert_cupy_array(a)
-    _util._assert_rank2(a)
 
     if mode not in ('reduced', 'complete', 'r', 'raw'):
         if mode in ('f', 'full', 'e', 'economic'):
             msg = 'The deprecated mode \'{}\' is not supported'.format(mode)
-            raise ValueError(msg)
         else:
-            raise ValueError('Unrecognized mode \'{}\''.format(mode))
+            msg = 'Unrecognized mode \'{}\''.format(mode)
+        raise ValueError(msg)
+    if a.ndim > 2:
+        cupy._util.experimental('cupy.linalg.qr')
+        return _qr_batched(a, mode)
 
     # support float32, float64, complex64, and complex128
     dtype, out_dtype = _util.linalg_common_type(a)
@@ -393,6 +437,7 @@ def _svd_batched(a, full_matrices, compute_uv):
         # copy (via possible type casting) is done in _gesvd_batched
         # note: _gesvd_batched returns V, not V^H
         out = _gesvd_batched(a, dtype.char, full_matrices, compute_uv, False)
+
     if compute_uv:
         u, s, v = out
         u = u.astype(uv_dtype, copy=False)
diff --git a/cupy/linalg/_util.py b/cupy/linalg/_util.py
--- a/cupy/linalg/_util.py
+++ b/cupy/linalg/_util.py
@@ -139,9 +139,10 @@ def _tril(x, k=0):
     return x
 
 
+# support a batch of matrices
 _triu_kernel = _core.ElementwiseKernel(
     'int64 k', 'S x',
-    'x = (_ind.get()[1] - _ind.get()[0] >= k) ? x : 0',
+    'x = (_ind.get()[_ind.ndim - 1] - _ind.get()[_ind.ndim - 2] >= k) ? x : 0',
     'cupy_triu_kernel',
     reduce_dims=False
 )


Test patch:
diff --git a/tests/cupy_tests/linalg_tests/test_decomposition.py b/tests/cupy_tests/linalg_tests/test_decomposition.py
--- a/tests/cupy_tests/linalg_tests/test_decomposition.py
+++ b/tests/cupy_tests/linalg_tests/test_decomposition.py
@@ -7,6 +7,8 @@
 from cupy._core.internal import prod
 from cupy import cusolver
 from cupy import testing
+from cupy.cuda import driver
+from cupy.cuda import runtime
 from cupy.testing import _condition
 import cupyx
 
@@ -101,11 +103,41 @@ def test_decomposition(self, dtype):
 class TestQRDecomposition(unittest.TestCase):
 
     @testing.for_dtypes('fdFD')
-    def check_mode(self, array, mode, dtype):
+    def check_mode(self, array, mode, dtype, batched=False):
+        if runtime.is_hip and driver.get_build_version() < 307:
+            if dtype in (numpy.complex64, numpy.complex128):
+                pytest.skip('ungqr unsupported')
+
         a_cpu = numpy.asarray(array, dtype=dtype)
         a_gpu = cupy.asarray(array, dtype=dtype)
-        result_cpu = numpy.linalg.qr(a_cpu, mode=mode)
         result_gpu = cupy.linalg.qr(a_gpu, mode=mode)
+
+        if ((not batched)
+                or (numpy.lib.NumpyVersion(numpy.__version__) >= '1.22.0')):
+            result_cpu = numpy.linalg.qr(a_cpu, mode=mode)
+            self._check_result(result_cpu, result_gpu)
+        else:
+            # We still want to test it to gain confidence...
+            # TODO(leofang): Use @testing.with_requires('numpy>=1.22') once
+            # NumPy 1.22 is out, and clean up this helper function
+            batch_shape = a_cpu.shape[:-2]
+            batch_size = prod(batch_shape)
+            a_cpu = a_cpu.reshape(batch_size, *a_cpu.shape[-2:])
+            for i in range(batch_size):
+                res_cpu = numpy.linalg.qr(a_cpu[i], mode=mode)
+                if isinstance(result_gpu, tuple):
+                    q_gpu, r_gpu = result_gpu
+                    q_gpu = q_gpu.reshape(batch_size, *q_gpu.shape[-2:])
+                    idx = -1 if mode == 'raw' else -2
+                    r_gpu = r_gpu.reshape(batch_size, *r_gpu.shape[idx:])
+                    res_gpu = (q_gpu[i], r_gpu[i])
+                    self._check_result(res_cpu, res_gpu)
+                else:  # mode == 'r'
+                    res_gpu = result_gpu.reshape(
+                        batch_size, *result_gpu.shape[-2:])[i]
+                    self._check_result(res_cpu, res_gpu)
+
+    def _check_result(self, result_cpu, result_gpu):
         if isinstance(result_cpu, tuple):
             for b_cpu, b_gpu in zip(result_cpu, result_gpu):
                 assert b_cpu.dtype == b_gpu.dtype
@@ -121,11 +153,43 @@ def test_mode(self):
         self.check_mode(numpy.random.randn(3, 3), mode=self.mode)
         self.check_mode(numpy.random.randn(5, 4), mode=self.mode)
 
+    @testing.fix_random()
+    @_condition.repeat(3, 10)
+    def test_mode_rank3(self):
+        self.check_mode(numpy.random.randn(3, 2, 4),
+                        mode=self.mode, batched=True)
+        self.check_mode(numpy.random.randn(4, 3, 3),
+                        mode=self.mode, batched=True)
+        self.check_mode(numpy.random.randn(2, 5, 4),
+                        mode=self.mode, batched=True)
+
+    @testing.fix_random()
+    @_condition.repeat(3, 10)
+    def test_mode_rank4(self):
+        self.check_mode(numpy.random.randn(2, 3, 2, 4),
+                        mode=self.mode, batched=True)
+        self.check_mode(numpy.random.randn(2, 4, 3, 3),
+                        mode=self.mode, batched=True)
+        self.check_mode(numpy.random.randn(2, 2, 5, 4),
+                        mode=self.mode, batched=True)
+
     @testing.with_requires('numpy>=1.16')
     def test_empty_array(self):
         self.check_mode(numpy.empty((0, 3)), mode=self.mode)
         self.check_mode(numpy.empty((3, 0)), mode=self.mode)
 
+    @testing.with_requires('numpy>=1.22')
+    def test_empty_array_rank3(self):
+        self.check_mode(numpy.empty((0, 3, 2)), mode=self.mode)
+        self.check_mode(numpy.empty((3, 0, 2)), mode=self.mode)
+        self.check_mode(numpy.empty((3, 2, 0)), mode=self.mode)
+        self.check_mode(numpy.empty((0, 3, 3)), mode=self.mode)
+        self.check_mode(numpy.empty((3, 0, 3)), mode=self.mode)
+        self.check_mode(numpy.empty((3, 3, 0)), mode=self.mode)
+        self.check_mode(numpy.empty((0, 2, 3)), mode=self.mode)
+        self.check_mode(numpy.empty((2, 0, 3)), mode=self.mode)
+        self.check_mode(numpy.empty((2, 3, 0)), mode=self.mode)
+
 
 @testing.parameterize(*testing.product({
     'full_matrices': [True, False],