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),
)