Add dpnp.linalg.lstsq() implementation

dpnp/linalg/dpnp_iface_linalg.py

@@ -42,13 +42,15 @@
 import dpnp
 
 from .dpnp_utils_linalg import (
+    check_2d,
     check_stacked_2d,
     check_stacked_square,
     dpnp_cholesky,
     dpnp_cond,
     dpnp_det,
     dpnp_eigh,
     dpnp_inv,
+    dpnp_lstsq,
     dpnp_matrix_power,
     dpnp_matrix_rank,
     dpnp_multi_dot,
@@ -69,6 +71,7 @@
     "eigvals",
     "eigvalsh",
     "inv",
+    "lstsq",
     "matrix_power",
     "matrix_rank",
     "multi_dot",
@@ -594,6 +597,76 @@ def inv(a):
     return dpnp_inv(a)
 
 
+def lstsq(a, b, rcond=None):
+    """
+    Return the least-squares solution to a linear matrix equation.
+
+    For full documentation refer to :obj:`numpy.linalg.lstsq`.
+
+    Parameters
+    ----------
+    a : (M, N) {dpnp.ndarray, usm_ndarray}
+        "Coefficient" matrix.
+    b : {(M,), (M, K)} {dpnp.ndarray, usm_ndarray}
+        Ordinate or "dependent variable" values.
+        If `b` is two-dimensional, the least-squares solution
+        is calculated for each of the `K` columns of `b`.
+    rcond : float, optional
+        Cut-off ratio for small singular values of `a`.
+        For the purposes of rank determination, singular values are treated
+        as zero if they are smaller than `rcond` times the largest singular
+        value of `a`.
+        The default uses the machine precision times ``max(M, N)``.  Passing
+        ``-1`` will use machine precision.
+
+    Returns
+    -------
+    x : {(N,), (N, K)} dpnp.ndarray
+        Least-squares solution. If `b` is two-dimensional,
+        the solutions are in the `K` columns of `x`.
+    residuals : {(1,), (K,), (0,)} dpnp.ndarray
+        Sums of squared residuals: Squared Euclidean 2-norm for each column in
+        ``b - a @ x``.
+        If the rank of `a` is < N or M <= N, this is an empty array.
+        If `b` is 1-dimensional, this is a (1,) shape array.
+        Otherwise the shape is (K,).
+    rank : int
+        Rank of matrix `a`.
+    s : (min(M, N),) dpnp.ndarray
+        Singular values of `a`.
+
+    Examples
+    --------
+    Fit a line, ``y = mx + c``, through some noisy data-points:
+    >>> import dpnp as np
+    >>> x = np.array([0, 1, 2, 3])
+    >>> y = np.array([-1, 0.2, 0.9, 2.1])
+
+    By examining the coefficients, we see that the line should have a
+    gradient of roughly 1 and cut the y-axis at, more or less, -1.
+
+    We can rewrite the line equation as ``y = Ap``, where ``A = [[x 1]]``
+    and ``p = [[m], [c]]``.  Now use `lstsq` to solve for `p`:
+
+    >>> A = np.vstack([x, np.ones(len(x))]).T
+    >>> A
+    array([[ 0.,  1.],
+           [ 1.,  1.],
+           [ 2.,  1.],
+           [ 3.,  1.]])
+
+    >>> m, c = np.linalg.lstsq(A, y, rcond=None)[0]
+    >>> m, c
+    (1.0 -0.95) # may vary
+
+    """
+
+    dpnp.check_supported_arrays_type(a, b)
+    check_2d(a)
+
+    return dpnp_lstsq(a, b, rcond=rcond)
+
+
 def matrix_power(a, n):
     """
     Raise a square matrix to the (integer) power `n`.

dpnp/linalg/dpnp_utils_linalg.py

@@ -42,6 +42,7 @@
     "dpnp_det",
     "dpnp_eigh",
     "dpnp_inv",
+    "dpnp_lstsq",
     "dpnp_matrix_power",
     "dpnp_matrix_rank",
     "dpnp_multi_dot",
@@ -653,6 +654,37 @@ def _triu_inplace(a, host_tasks, depends=None):
     return out
 
 
+def check_2d(*arrays):
+    """
+    Return ``True`` if each array in `arrays` is exactly two dimensions.
+
+    If any array is not two-dimensional, `dpnp.linalg.LinAlgError` will be raised.
+
+    Parameters
+    ----------
+    arrays : {dpnp.ndarray, usm_ndarray}
+        A sequence of input arrays to check for dimensionality.
+
+    Returns
+    -------
+    out : bool
+        ``True`` if each array in `arrays` is exactly two-dimensional.
+
+    Raises
+    ------
+    dpnp.linalg.LinAlgError
+        If any array in `arrays` is not exactly two-dimensional.
+
+    """
+
+    for a in arrays:
+        if a.ndim != 2:
+            raise dpnp.linalg.LinAlgError(
+                f"{a.ndim}-dimensional array given. The input "
+                "array must be exactly two-dimensional"
+            )
+
+
 def check_stacked_2d(*arrays):
     """
     Return ``True`` if each array in `arrays` has at least two dimensions.
@@ -1224,6 +1256,77 @@ def dpnp_inv(a):
     return b_f
 
 
+def _nrm2_last_axis(x):
+    real_dtype = _real_type(x.dtype)
+    x = dpnp.ascontiguousarray(x)
+    if dpnp.issubdtype(x.dtype, dpnp.complexfloating):
+        squared_abs = dpnp.sum(dpnp.abs(x) ** 2, axis=-1, dtype=real_dtype)
+        return squared_abs
+    else:
+        return dpnp.sum(dpnp.square(x), axis=-1, dtype=real_dtype)
+
+
+def dpnp_lstsq(a, b, rcond=None):
+    """
+    dpnp_lstsq(a, b, rcond=None)
+
+    Return the least-squares solution to a linear matrix equation.
+
+    """
+
+    if b.ndim > 2:
+        raise dpnp.linalg.LinAlgError(
+            f"{b.ndim}-dimensional array given. The input "
+            "array must be exactly two-dimensional"
+        )
+
+    m, n = a.shape[-2:]
+    m2 = b.shape[0]
+    if m != m2:
+        raise dpnp.linalg.LinAlgError("Incompatible dimensions")
+
+    u, s, vh = dpnp_svd(a, full_matrices=False)
+
+    if rcond is None:
+        rcond = dpnp.finfo(s.dtype).eps * max(m, n)
+    elif rcond <= 0 or rcond >= 1:
+        # some doc of gelss/gelsd says "rcond < 0", but it's not true!
+        rcond = dpnp.finfo(s.dtype).eps
+
+    res_usm_type, exec_q = get_usm_allocations([a, b])
+
+    # number of singular values and matrix rank
+    s1 = 1 / s
+    rank = dpnp.array(
+        s.size, dtype="int32", sycl_queue=exec_q, usm_type=res_usm_type
+    )
+    if s.size > 0:
+        cutoff = rcond * s.max()
+        sing_vals = s <= cutoff
+        s1[sing_vals] = 0
+        rank -= sing_vals.sum(dtype="int32")
+
+    # Solve the least-squares solution
+    # x = vh.T.conj() @ diag(s1) @ u.T.conj() @ b
+    z = (dpnp.dot(b.T, u.conj()) * s1).T
+    x = dpnp.dot(vh.T.conj(), z)
+    # Calculate squared Euclidean 2-norm for each column in b - a*x
+    if m <= n or rank != n:
+        resids = dpnp.empty(
+            (0,), dtype=s.dtype, sycl_queue=exec_q, usm_type=res_usm_type
+        )
+    else:
+        e = b - a.dot(x)
+        resids = dpnp.atleast_1d(_nrm2_last_axis(e.T))
+
+    # We need to copy singular values array to USM memory according
+    # to compute follows data if its type of SYCL USM allocation
+    # does not match with res_usm_type
+    if s.usm_type != res_usm_type:
+        s = dpnp.copy(a, sycl_queue=exec_q, usm_type=res_usm_type)
+    return x, resids, rank, s
+
+
 def dpnp_matrix_power(a, n):
     """
     dpnp_matrix_power(a, n)

tests/test_sycl_queue.py

@@ -2022,3 +2022,35 @@ def test_tensorsolve(device):
     result_queue = result.sycl_queue
 
     assert_sycl_queue_equal(result_queue, a_dp.sycl_queue)
+
+
+@pytest.mark.parametrize(
+    "device",
+    valid_devices,
+    ids=[device.filter_string for device in valid_devices],
+)
+@pytest.mark.parametrize(
+    ["m", "n", "nrhs"],
+    [
+        (4, 2, 2),
+        (4, 0, 1),
+        (4, 2, 0),
+        (0, 0, 0),
+    ],
+)
+def test_lstsq(m, n, nrhs, device):
+    a_np = numpy.arange(m * n).reshape(m, n).astype(dpnp.default_float_type())
+    b_np = numpy.ones((m, nrhs)).astype(dpnp.default_float_type())
+
+    a_dp = dpnp.array(a_np, device=device)
+    b_dp = dpnp.array(b_np, device=device)
+
+    result_dp = dpnp.linalg.lstsq(a_dp, b_dp)
+    result = numpy.linalg.lstsq(a_np, b_np)
+
+    for param_dp, param_np in zip(result_dp, result):
+        assert_dtype_allclose(param_dp, param_np)
+
+    for param_dp in result_dp:
+        assert_sycl_queue_equal(param_dp.sycl_queue, a_dp.sycl_queue)
+        assert_sycl_queue_equal(param_dp.sycl_queue, b_dp.sycl_queue)

tests/test_usm_type.py

@@ -1170,3 +1170,28 @@ def test_tensorsolve(usm_type_a, usm_type_b):
     assert a.usm_type == usm_type_a
     assert b.usm_type == usm_type_b
     assert result.usm_type == du.get_coerced_usm_type([usm_type_a, usm_type_b])
+
+
+@pytest.mark.parametrize("usm_type_a", list_of_usm_types, ids=list_of_usm_types)
+@pytest.mark.parametrize("usm_type_b", list_of_usm_types, ids=list_of_usm_types)
+@pytest.mark.parametrize(
+    ["m", "n", "nrhs"],
+    [
+        (4, 2, 2),
+        (4, 0, 1),
+        (4, 2, 0),
+        (0, 0, 0),
+    ],
+)
+def test_lstsq(m, n, nrhs, usm_type_a, usm_type_b):
+    a = dp.arange(m * n, usm_type=usm_type_a).reshape(m, n)
+    b = dp.ones((m, nrhs), usm_type=usm_type_b)
+
+    result = dp.linalg.lstsq(a, b)
+
+    assert a.usm_type == usm_type_a
+    assert b.usm_type == usm_type_b
+    for param in result:
+        assert param.usm_type == du.get_coerced_usm_type(
+            [usm_type_a, usm_type_b]
+        )

tests/third_party/cupy/linalg_tests/test_solve.py

@@ -226,6 +226,98 @@ def test_pinv_size_0(self):
         self.check_x((2, 0, 3), rcond=1e-15)
 
 
+class TestLstsq:
+    @testing.for_dtypes("ifdFD")
+    @testing.numpy_cupy_allclose(atol=1e-3, type_check=has_support_aspect64())
+    def check_lstsq_solution(
+        self, a_shape, b_shape, seed, rcond, xp, dtype, singular=False
+    ):
+        if singular:
+            m, n = a_shape
+            rank = min(m, n) - 1
+            a = xp.matmul(
+                testing.shaped_random(
+                    (m, rank), xp, dtype=dtype, scale=3, seed=seed
+                ),
+                testing.shaped_random(
+                    (rank, n), xp, dtype=dtype, scale=3, seed=seed + 42
+                ),
+            )
+        else:
+            a = testing.shaped_random(a_shape, xp, dtype=dtype, seed=seed)
+        b = testing.shaped_random(b_shape, xp, dtype=dtype, seed=seed + 37)
+        a_copy = a.copy()
+        b_copy = b.copy()
+        results = xp.linalg.lstsq(a, b, rcond)
+        if xp is cupy:
+            testing.assert_array_equal(a_copy, a)
+            testing.assert_array_equal(b_copy, b)
+        return results
+
+    def check_invalid_shapes(self, a_shape, b_shape):
+        a = cupy.random.rand(*a_shape)
+        b = cupy.random.rand(*b_shape)
+        with pytest.raises(cupy.linalg.LinAlgError):
+            cupy.linalg.lstsq(a, b, rcond=None)
+
+    def test_lstsq_solutions(self):
+        # Comapres numpy.linalg.lstsq and cupy.linalg.lstsq solutions for:
+        #   a shapes range from (3, 3) to (5, 3) and (3, 5)
+        #   b shapes range from (i, 3) to (i, )
+        #   sets a random seed for deterministic testing
+        for i in range(3, 6):
+            for j in range(3, 6):
+                for k in range(2, 4):
+                    seed = i + j + k
+                    # check when b has shape (i, k)
+                    self.check_lstsq_solution((i, j), (i, k), seed, rcond=-1)
+                    self.check_lstsq_solution((i, j), (i, k), seed, rcond=None)
+                    self.check_lstsq_solution((i, j), (i, k), seed, rcond=0.5)
+                    self.check_lstsq_solution(
+                        (i, j), (i, k), seed, rcond=1e-6, singular=True
+                    )
+                # check when b has shape (i, )
+                self.check_lstsq_solution((i, j), (i,), seed + 1, rcond=-1)
+                self.check_lstsq_solution((i, j), (i,), seed + 1, rcond=None)
+                self.check_lstsq_solution((i, j), (i,), seed + 1, rcond=0.5)
+                self.check_lstsq_solution(
+                    (i, j), (i,), seed + 1, rcond=1e-6, singular=True
+                )
+
+    @pytest.mark.parametrize("rcond", [-1, None, 0.5])
+    @pytest.mark.parametrize("k", [None, 0, 1, 4])
+    @pytest.mark.parametrize(("i", "j"), [(0, 0), (3, 0), (0, 7)])
+    @testing.for_dtypes("ifdFD")
+    @testing.numpy_cupy_allclose(rtol=1e-6, type_check=has_support_aspect64())
+    def test_lstsq_empty_matrix(self, xp, dtype, i, j, k, rcond):
+        a = xp.empty((i, j), dtype=dtype)
+        assert a.size == 0
+        b_shape = (i,) if k is None else (i, k)
+        b = testing.shaped_random(b_shape, xp, dtype)
+        return xp.linalg.lstsq(a, b, rcond)
+
+    def test_invalid_shapes(self):
+        self.check_invalid_shapes((4, 3), (3,))
+        self.check_invalid_shapes((3, 3, 3), (2, 2))
+        self.check_invalid_shapes((3, 3, 3), (3, 3))
+        self.check_invalid_shapes((3, 3), (3, 3, 3))
+        self.check_invalid_shapes((2, 2), (10,))
+        self.check_invalid_shapes((3, 3), (2, 2))
+        self.check_invalid_shapes((4, 3), (10, 3, 3))
+
+    # dpnp.linalg.lstsq() does not raise a FutureWarning
+    # because the next version of numpy will not raise it
+    # and this test will be removed
+    @pytest.mark.skip("Not implemented")
+    @testing.for_float_dtypes(no_float16=True)
+    @testing.numpy_cupy_allclose(atol=1e-3)
+    def test_warn_rcond(self, xp, dtype):
+        a = testing.shaped_random((3, 3), xp, dtype)
+        b = testing.shaped_random((3,), xp, dtype)
+        with testing.assert_warns(FutureWarning):
+            return xp.linalg.lstsq(a, b)
+
+
 class TestTensorInv(unittest.TestCase):
     @testing.for_dtypes("ifdFD")
     @_condition.retry(10)