Repository: cupy/cupy

Issue:
cupy.linalg.{svd, pinv} should support broadcasting
Problem
-----------
The NumPy implementation of [pinv](https://docs.scipy.org/doc/numpy/reference/generated/numpy.linalg.pinv.html) supports broadcasting, but the CuPy implementation does not.

Conditions
--------------
```
CuPy Version          : 7.0.0
CUDA Root             : /local/cuda-10.0
CUDA Build Version    : 10010
CUDA Driver Version   : 10020
CUDA Runtime Version  : 10010
cuBLAS Version        : 10201
cuFFT Version         : 10101
cuRAND Version        : 10101
cuSOLVER Version      : (10, 2, 0)
cuSPARSE Version      : 10300
NVRTC Version         : (10, 1)
cuDNN Build Version   : 7604
cuDNN Version         : 7605
NCCL Build Version    : 2406
NCCL Runtime Version  : 2506
```
Reproduce
--------------
```python
import numpy as np
import cupy as cp

A = np.random.rand(3, 4, 5)
xA = cp.array(A)
np.linalg.pinv(A)
cp.linalg.pinv(xA)
```
```
Traceback (most recent call last):
  File "pinv.py", line 7, in <module>
    cp.linalg.pinv(xA)
  File "/home/beams/DCHING/miniconda3/envs/tike/lib/python3.7/site-packages/cupy/linalg/solve.py", line 421, in pinv
    u, s, vt = decomposition.svd(a.conj(), full_matrices=False)
  File "/home/beams/DCHING/miniconda3/envs/tike/lib/python3.7/site-packages/cupy/linalg/decomposition.py", line 245, in svd
    util._assert_rank2(a)
  File "/home/beams/DCHING/miniconda3/envs/tike/lib/python3.7/site-packages/cupy/linalg/util.py", line 20, in _assert_rank2
    'two-dimensional'.format(a.ndim))
numpy.linalg.LinAlgError: 3-dimensional array given. Array must be two-dimensional
```
Solution
-----------
Change CuPy implementation to allow broadcasting on N+2 dimensional arrays. This may also require updating `cupy.linalg.svd` because `pinv` relies on `svd` and NumPy allows broadcasting for that function too.





Discussion:
Broadcasting has been enabled for some `linalg` functions since [`numpy=1.8`](https://numpy.org/doc/1.18/reference/routines.linalg.html#module-numpy.linalg)
This problem is being addressed in #2337
#2337 implements efficient batched SVD of **small** matrices.  The first support could be done by a trivial loop (in Python or in Cython).
Related to #3470.
Related: https://github.com/data-apis/array-api/pull/114, https://github.com/data-apis/array-api/pull/118 

Resolving patch:
diff --git a/cupy/linalg/_decomposition.py b/cupy/linalg/_decomposition.py
--- a/cupy/linalg/_decomposition.py
+++ b/cupy/linalg/_decomposition.py
@@ -412,6 +412,7 @@ def _svd_batched(a, a_dtype, full_matrices, compute_uv):
         return s
 
 
+# TODO(leofang): support the hermitian keyword?
 def svd(a, full_matrices=True, compute_uv=True):
     """Singular Value Decomposition.
 
diff --git a/cupy/linalg/_solve.py b/cupy/linalg/_solve.py
--- a/cupy/linalg/_solve.py
+++ b/cupy/linalg/_solve.py
@@ -382,6 +382,7 @@ def _batched_inv(a):
     return c.reshape(a_shape)
 
 
+# TODO(leofang): support the hermitian keyword?
 def pinv(a, rcond=1e-15):
     """Compute the Moore-Penrose pseudoinverse of a matrix.
 
@@ -390,13 +391,15 @@ def pinv(a, rcond=1e-15):
     Note that it automatically removes small singular values for stability.
 
     Args:
-        a (cupy.ndarray): The matrix with dimension ``(M, N)``
-        rcond (float): Cutoff parameter for small singular values.
-            For stability it computes the largest singular value denoted by
-            ``s``, and sets all singular values smaller than ``s`` to zero.
+        a (cupy.ndarray): The matrix with dimension ``(..., M, N)``
+        rcond (float or cupy.ndarray): Cutoff parameter for small singular
+            values. For stability it computes the largest singular value
+            denoted by ``s``, and sets all singular values smaller than
+            ``rcond * s`` to zero. Broadcasts against the stack of matrices.
 
     Returns:
-        cupy.ndarray: The pseudoinverse of ``a`` with dimension ``(N, M)``.
+        cupy.ndarray: The pseudoinverse of ``a`` with dimension
+            ``(..., N, M)``.
 
     .. warning::
         This function calls one or more cuSOLVER routine(s) which may yield
@@ -407,11 +410,18 @@ def pinv(a, rcond=1e-15):
 
     .. seealso:: :func:`numpy.linalg.pinv`
     """
+    if a.size == 0:
+        m, n = a.shape[-2:]
+        return cupy.empty(a.shape[:-2] + (n, m), dtype=a.dtype)
     u, s, vt = _decomposition.svd(a.conj(), full_matrices=False)
-    cutoff = rcond * s.max()
-    s1 = 1 / s
-    s1[s <= cutoff] = 0
-    return cupy.dot(vt.T, s1[:, None] * u.T)
+
+    # discard small singular values
+    cutoff = rcond * cupy.amax(s, axis=-1)
+    leq = s <= cutoff[..., None]
+    cupy.reciprocal(s, out=s)
+    s[leq] = 0
+
+    return cupy.matmul(vt.swapaxes(-2, -1), s[..., None] * u.swapaxes(-2, -1))
 
 
 def tensorinv(a, ind=2):


Test patch:
diff --git a/tests/cupy_tests/linalg_tests/test_solve.py b/tests/cupy_tests/linalg_tests/test_solve.py
--- a/tests/cupy_tests/linalg_tests/test_solve.py
+++ b/tests/cupy_tests/linalg_tests/test_solve.py
@@ -156,18 +156,17 @@ def check_x(self, a_shape, rcond, dtype):
         a_gpu = testing.shaped_random(a_shape, dtype=dtype)
         a_cpu = cupy.asnumpy(a_gpu)
         a_gpu_copy = a_gpu.copy()
+        if not isinstance(rcond, float):
+            rcond = numpy.asarray(rcond)
         result_cpu = numpy.linalg.pinv(a_cpu, rcond=rcond)
+        if not isinstance(rcond, float):
+            rcond = cupy.asarray(rcond)
         result_gpu = cupy.linalg.pinv(a_gpu, rcond=rcond)
 
         assert result_cpu.dtype == result_gpu.dtype
         cupy.testing.assert_allclose(result_cpu, result_gpu, atol=1e-3)
         cupy.testing.assert_array_equal(a_gpu_copy, a_gpu)
 
-    def check_shape(self, a_shape):
-        a = cupy.random.rand(*a_shape)
-        with self.assertRaises(ValueError):
-            cupy.linalg.pinv(a)
-
     def test_pinv(self):
         self.check_x((3, 3), rcond=1e-15)
         self.check_x((2, 4), rcond=1e-15)
@@ -177,11 +176,22 @@ def test_pinv(self):
         self.check_x((2, 5), rcond=0.5)
         self.check_x((5, 3), rcond=0.6)
 
-    def test_invalid_shape(self):
-        # TODO(leofang): NumPy supports batched pinv. Since we now support
-        # batched svd (which is behind pinv), this should not raise
-        self.check_shape((2, 3, 4))
-        self.check_shape((4, 3, 2, 1))
+    def test_pinv_batched(self):
+        self.check_x((2, 3, 4), rcond=1e-15)
+        self.check_x((2, 3, 4, 5), rcond=1e-15)
+
+    def test_pinv_batched_vector_rcond(self):
+        self.check_x((2, 3, 4), rcond=[0.2, 0.8])
+        self.check_x((2, 3, 4, 5),
+                     rcond=[[0.2, 0.9, 0.1],
+                            [0.7, 0.2, 0.5]])
+
+    def test_pinv_size_0(self):
+        self.check_x((3, 0), rcond=1e-15)
+        self.check_x((0, 3), rcond=1e-15)
+        self.check_x((0, 0), rcond=1e-15)
+        self.check_x((0, 2, 3), rcond=1e-15)
+        self.check_x((2, 0, 3), rcond=1e-15)
 
 
 @testing.gpu