def _solve(a, b): from cupy.linalg._util import (_assert_stacked_2d, _assert_stacked_square, linalg_common_type) _commonType = functools.partial(linalg_common_type, reject_float16=False) _assert_stacked_2d(a) _assert_stacked_square(a) _, result_t = _commonType(a, b) if b.ndim == 1: # (M,) -> (M, 1) old_shape = b.shape b = b.reshape(-1, 1) elif b.ndim == a.ndim - 1: # x1 has shape (M, M) or (..., M, M) raise ValueError('x2 must have shape (M,) or (..., M, K); ' '(..., M) is not allowed') else: # (..., M, K) => no change old_shape = None r = np.linalg.solve(a, b).astype(result_t, copy=False) r = r.reshape(old_shape) if old_shape else r return r
def eigvalsh(a, UPLO='L'): """ Compute the eigenvalues of a complex Hermitian or real symmetric matrix. Main difference from eigh: the eigenvectors are not computed. Args: a (cupy.ndarray): A symmetric 2-D square matrix ``(M, M)`` or a batch of symmetric 2-D square matrices ``(..., M, M)``. UPLO (str): Select from ``'L'`` or ``'U'``. It specifies which part of ``a`` is used. ``'L'`` uses the lower triangular part of ``a``, and ``'U'`` uses the upper triangular part of ``a``. Returns: cupy.ndarray: Returns eigenvalues as a vector ``w``. For batch input, ``w[k]`` is a vector of eigenvalues of matrix ``a[k]``. .. warning:: This function calls one or more cuSOLVER routine(s) which may yield invalid results if input conditions are not met. To detect these invalid results, you can set the `linalg` configuration to a value that is not `ignore` in :func:`cupyx.errstate` or :func:`cupyx.seterr`. .. seealso:: :func:`numpy.linalg.eigvalsh` """ _util._assert_stacked_2d(a) _util._assert_stacked_square(a) if a.ndim > 2 or runtime.is_hip: return cupy.cusolver.syevj(a, UPLO, False) else: return _syevd(a, UPLO, False)[0]
def solve(a, b): """Solves a linear matrix equation. It computes the exact solution of ``x`` in ``ax = b``, where ``a`` is a square and full rank matrix. Args: a (cupy.ndarray): The matrix with dimension ``(..., M, M)``. b (cupy.ndarray): The matrix with dimension ``(...,M)`` or ``(..., M, K)``. Returns: cupy.ndarray: The matrix with dimension ``(..., M)`` or ``(..., M, K)``. .. warning:: This function calls one or more cuSOLVER routine(s) which may yield invalid results if input conditions are not met. To detect these invalid results, you can set the `linalg` configuration to a value that is not `ignore` in :func:`cupyx.errstate` or :func:`cupyx.seterr`. .. seealso:: :func:`numpy.linalg.solve` """ if a.ndim > 2 and a.shape[-1] <= get_batched_gesv_limit(): # Note: There is a low performance issue in batched_gesv when matrix is # large, so it is not used in such cases. return batched_gesv(a, b) # TODO(kataoka): Move the checks to the beginning _util._assert_cupy_array(a, b) _util._assert_stacked_2d(a) _util._assert_stacked_square(a) if not ((a.ndim == b.ndim or a.ndim == b.ndim + 1) and a.shape[:-1] == b.shape[:a.ndim - 1]): raise ValueError( 'a must have (..., M, M) shape and b must have (..., M) ' 'or (..., M, K)') dtype, out_dtype = _util.linalg_common_type(a, b) if a.ndim == 2: # prevent 'a' and 'b' to be overwritten a = a.astype(dtype, copy=True, order='F') b = b.astype(dtype, copy=True, order='F') cupyx.lapack.gesv(a, b) return b.astype(out_dtype, copy=False) # prevent 'a' to be overwritten a = a.astype(dtype, copy=True, order='C') x = cupy.empty_like(b, dtype=out_dtype) shape = a.shape[:-2] for i in range(numpy.prod(shape)): index = numpy.unravel_index(i, shape) # prevent 'b' to be overwritten bi = b[index].astype(dtype, copy=True, order='F') cupyx.lapack.gesv(a[index], bi) x[index] = bi return x
def cholesky(a): """Cholesky decomposition. Decompose a given two-dimensional square matrix into ``L * L.T``, where ``L`` is a lower-triangular matrix and ``.T`` is a conjugate transpose operator. Args: a (cupy.ndarray): The input matrix with dimension ``(N, N)`` Returns: cupy.ndarray: The lower-triangular matrix. .. warning:: This function calls one or more cuSOLVER routine(s) which may yield invalid results if input conditions are not met. To detect these invalid results, you can set the `linalg` configuration to a value that is not `ignore` in :func:`cupyx.errstate` or :func:`cupyx.seterr`. .. seealso:: :func:`numpy.linalg.cholesky` """ _util._assert_cupy_array(a) _util._assert_stacked_2d(a) _util._assert_stacked_square(a) if a.ndim > 2: return _potrf_batched(a) dtype, out_dtype = _util.linalg_common_type(a) x = a.astype(dtype, order='C', copy=True) n = len(a) handle = device.get_cusolver_handle() dev_info = cupy.empty(1, dtype=numpy.int32) if dtype == 'f': potrf = cusolver.spotrf potrf_bufferSize = cusolver.spotrf_bufferSize elif dtype == 'd': potrf = cusolver.dpotrf potrf_bufferSize = cusolver.dpotrf_bufferSize elif dtype == 'F': potrf = cusolver.cpotrf potrf_bufferSize = cusolver.cpotrf_bufferSize else: # dtype == 'D': potrf = cusolver.zpotrf potrf_bufferSize = cusolver.zpotrf_bufferSize buffersize = potrf_bufferSize(handle, cublas.CUBLAS_FILL_MODE_UPPER, n, x.data.ptr, n) workspace = cupy.empty(buffersize, dtype=dtype) potrf(handle, cublas.CUBLAS_FILL_MODE_UPPER, n, x.data.ptr, n, workspace.data.ptr, buffersize, dev_info.data.ptr) cupy.linalg._util._check_cusolver_dev_info_if_synchronization_allowed( potrf, dev_info) _util._tril(x, k=0) return x.astype(out_dtype, copy=False)
def eigh(a, UPLO='L'): """ Return the eigenvalues and eigenvectors of a complex Hermitian (conjugate symmetric) or a real symmetric matrix. Returns two objects, a 1-D array containing the eigenvalues of `a`, and a 2-D square array or matrix (depending on the input type) of the corresponding eigenvectors (in columns). Args: a (cupy.ndarray): A symmetric 2-D square matrix ``(M, M)`` or a batch of symmetric 2-D square matrices ``(..., M, M)``. UPLO (str): Select from ``'L'`` or ``'U'``. It specifies which part of ``a`` is used. ``'L'`` uses the lower triangular part of ``a``, and ``'U'`` uses the upper triangular part of ``a``. Returns: tuple of :class:`~cupy.ndarray`: Returns a tuple ``(w, v)``. ``w`` contains eigenvalues and ``v`` contains eigenvectors. ``v[:, i]`` is an eigenvector corresponding to an eigenvalue ``w[i]``. For batch input, ``v[k, :, i]`` is an eigenvector corresponding to an eigenvalue ``w[k, i]`` of ``a[k]``. .. warning:: This function calls one or more cuSOLVER routine(s) which may yield invalid results if input conditions are not met. To detect these invalid results, you can set the `linalg` configuration to a value that is not `ignore` in :func:`cupyx.errstate` or :func:`cupyx.seterr`. .. seealso:: :func:`numpy.linalg.eigh` """ _util._assert_stacked_2d(a) _util._assert_stacked_square(a) if a.size == 0: _, v_dtype = _util.linalg_common_type(a) w_dtype = v_dtype.char.lower() w = cupy.empty(a.shape[:-1], w_dtype) v = cupy.empty(a.shape, v_dtype) return w, v if a.ndim > 2 or runtime.is_hip: w, v = cupy.cusolver.syevj(a, UPLO, True) return w, v else: return _syevd(a, UPLO, True)
def matrix_power(M, n): """Raise a square matrix to the (integer) power `n`. Args: M (~cupy.ndarray): Matrix to raise by power n. n (~int): Power to raise matrix to. Returns: ~cupy.ndarray: Output array. ..seealso:: :func:`numpy.linalg.matrix_power` """ _util._assert_cupy_array(M) _util._assert_stacked_2d(M) _util._assert_stacked_square(M) if not isinstance(n, int): raise TypeError('exponent must be an integer') if n == 0: return _util.stacked_identity_like(M) elif n < 0: M = _solve.inv(M) n *= -1 # short-cuts if n <= 3: if n == 1: return M elif n == 2: return cupy.matmul(M, M) else: return cupy.matmul(cupy.matmul(M, M), M) # binary decomposition to reduce the number of Matrix # multiplications for n > 3. result, Z = None, None for b in cupy.binary_repr(n)[::-1]: Z = M if Z is None else cupy.matmul(Z, Z) if b == '1': result = Z if result is None else cupy.matmul(result, Z) return result
def inv(a): """Computes the inverse of a matrix. This function computes matrix ``a_inv`` from n-dimensional regular matrix ``a`` such that ``dot(a, a_inv) == eye(n)``. Args: a (cupy.ndarray): The regular matrix Returns: cupy.ndarray: The inverse of a matrix. .. warning:: This function calls one or more cuSOLVER routine(s) which may yield invalid results if input conditions are not met. To detect these invalid results, you can set the `linalg` configuration to a value that is not `ignore` in :func:`cupyx.errstate` or :func:`cupyx.seterr`. .. seealso:: :func:`numpy.linalg.inv` """ _util._assert_cupy_array(a) _util._assert_stacked_2d(a) _util._assert_stacked_square(a) if a.ndim >= 3: return _batched_inv(a) dtype, out_dtype = _util.linalg_common_type(a) if a.size == 0: return cupy.empty(a.shape, out_dtype) order = 'F' if a._f_contiguous else 'C' # prevent 'a' to be overwritten a = a.astype(dtype, copy=True, order=order) b = cupy.eye(a.shape[0], dtype=dtype, order=order) if order == 'F': cupyx.lapack.gesv(a, b) else: cupyx.lapack.gesv(a.T, b.T) return b.astype(out_dtype, copy=False)
def batched_gesv(a, b): """Solves multiple linear matrix equations using cublas<t>getr[fs]Batched(). Computes the solution to system of linear equation ``ax = b``. Args: a (cupy.ndarray): The matrix with dimension ``(..., M, M)``. b (cupy.ndarray): The matrix with dimension ``(..., M)`` or ``(..., M, K)``. Returns: cupy.ndarray: The matrix with dimension ``(..., M)`` or ``(..., M, K)``. """ _util._assert_cupy_array(a, b) _util._assert_stacked_2d(a) _util._assert_stacked_square(a) # TODO(kataoka): Support broadcast if not ( (a.ndim == b.ndim or a.ndim == b.ndim + 1) and a.shape[:-1] == b.shape[:a.ndim - 1] ): raise ValueError( 'a must have (..., M, M) shape and b must have (..., M) ' 'or (..., M, K)') dtype, out_dtype = _util.linalg_common_type(a, b) if b.size == 0: return cupy.empty(b.shape, out_dtype) if dtype == 'f': t = 's' elif dtype == 'd': t = 'd' elif dtype == 'F': t = 'c' elif dtype == 'D': t = 'z' else: raise TypeError('invalid dtype') getrf = getattr(cublas, t + 'getrfBatched') getrs = getattr(cublas, t + 'getrsBatched') bs = numpy.prod(a.shape[:-2]) if a.ndim > 2 else 1 n = a.shape[-1] nrhs = b.shape[-1] if a.ndim == b.ndim else 1 b_shape = b.shape a_data_ptr = a.data.ptr b_data_ptr = b.data.ptr a = cupy.ascontiguousarray(a.reshape(bs, n, n).transpose(0, 2, 1), dtype=dtype) b = cupy.ascontiguousarray(b.reshape(bs, n, nrhs).transpose(0, 2, 1), dtype=dtype) if a.data.ptr == a_data_ptr: a = a.copy() if b.data.ptr == b_data_ptr: b = b.copy() if n > get_batched_gesv_limit(): warnings.warn('The matrix size ({}) exceeds the set limit ({})'. format(n, get_batched_gesv_limit())) handle = device.get_cublas_handle() lda = n a_step = lda * n * a.itemsize a_array = cupy.arange(a.data.ptr, a.data.ptr + a_step * bs, a_step, dtype=cupy.uintp) ldb = n b_step = ldb * nrhs * b.itemsize b_array = cupy.arange(b.data.ptr, b.data.ptr + b_step * bs, b_step, dtype=cupy.uintp) pivot = cupy.empty((bs, n), dtype=numpy.int32) dinfo = cupy.empty((bs,), dtype=numpy.int32) info = numpy.empty((1,), dtype=numpy.int32) # LU factorization (A = L * U) getrf(handle, n, a_array.data.ptr, lda, pivot.data.ptr, dinfo.data.ptr, bs) _util._check_cublas_info_array_if_synchronization_allowed(getrf, dinfo) # Solves Ax = b getrs(handle, cublas.CUBLAS_OP_N, n, nrhs, a_array.data.ptr, lda, pivot.data.ptr, b_array.data.ptr, ldb, info.ctypes.data, bs) if info[0] != 0: msg = 'Error reported by {} in cuBLAS. '.format(getrs.__name__) if info[0] < 0: msg += 'The {}-th parameter had an illegal value.'.format(-info[0]) raise linalg.LinAlgError(msg) return b.transpose(0, 2, 1).reshape(b_shape).astype(out_dtype, copy=False)
def slogdet(a): """Returns sign and logarithm of the determinant of an array. It calculates the natural logarithm of the determinant of a given value. Args: a (cupy.ndarray): The input matrix with dimension ``(..., N, N)``. Returns: tuple of :class:`~cupy.ndarray`: It returns a tuple ``(sign, logdet)``. ``sign`` represents each sign of the determinant as a real number ``0``, ``1`` or ``-1``. 'logdet' represents the natural logarithm of the absolute of the determinant. If the determinant is zero, ``sign`` will be ``0`` and ``logdet`` will be ``-inf``. The shapes of both ``sign`` and ``logdet`` are equal to ``a.shape[:-2]``. .. warning:: This function calls one or more cuSOLVER routine(s) which may yield invalid results if input conditions are not met. To detect these invalid results, you can set the `linalg` configuration to a value that is not `ignore` in :func:`cupyx.errstate` or :func:`cupyx.seterr`. .. warning:: To produce the same results as :func:`numpy.linalg.slogdet` for singular inputs, set the `linalg` configuration to `raise`. .. seealso:: :func:`numpy.linalg.slogdet` """ _util._assert_stacked_2d(a) _util._assert_stacked_square(a) dtype, sign_dtype = _util.linalg_common_type(a) logdet_dtype = numpy.dtype(sign_dtype.char.lower()) a_shape = a.shape shape = a_shape[:-2] n = a_shape[-2] if a.size == 0: # empty batch (result is empty, too) or empty matrices det([[]]) == 1 sign = cupy.ones(shape, sign_dtype) logdet = cupy.zeros(shape, logdet_dtype) return sign, logdet lu, ipiv, dev_info = _decomposition._lu_factor(a, dtype) # dev_info < 0 means illegal value (in dimensions, strides, and etc.) that # should never happen even if the matrix contains nan or inf. # TODO(kataoka): assert dev_info >= 0 if synchronization is allowed for # debugging purposes. diag = cupy.diagonal(lu, axis1=-2, axis2=-1) logdet = cupy.log(cupy.abs(diag)).sum(axis=-1) # ipiv is 1-origin non_zero = cupy.count_nonzero(ipiv != cupy.arange(1, n + 1), axis=-1) if dtype.kind == "f": non_zero += cupy.count_nonzero(diag < 0, axis=-1) # Note: sign == -1 ** (non_zero % 2) sign = (non_zero % 2) * -2 + 1 if dtype.kind == "c": sign = sign * cupy.prod(diag / cupy.abs(diag), axis=-1) sign = sign.astype(dtype) logdet = logdet.astype(logdet_dtype, copy=False) singular = dev_info > 0 return ( cupy.where(singular, sign_dtype.type(0), sign).reshape(shape), cupy.where(singular, logdet_dtype.type('-inf'), logdet).reshape(shape), )