def _syevj_batched(a, UPLO, with_eigen_vector):
if a.dtype == 'f' or a.dtype == 'e':
dtype = 'f'
inp_w_dtype = 'f'
inp_v_dtype = 'f'
ret_w_dtype = a.dtype
ret_v_dtype = a.dtype
elif a.dtype == 'd':
dtype = 'd'
inp_w_dtype = 'd'
inp_v_dtype = 'd'
ret_w_dtype = 'd'
ret_v_dtype = 'd'
elif a.dtype == 'F':
dtype = 'F'
inp_w_dtype = 'f'
inp_v_dtype = 'F'
ret_w_dtype = 'f'
ret_v_dtype = 'F'
elif a.dtype == 'D':
dtype = 'D'
inp_w_dtype = 'd'
inp_v_dtype = 'D'
ret_w_dtype = 'd'
ret_v_dtype = 'D'
else:
# NumPy uses float64 when an input is not floating point number.
dtype = 'd'
inp_w_dtype = 'd'
inp_v_dtype = 'd'
ret_w_dtype = 'd'
ret_v_dtype = 'd'
*batch_shape, m, lda = a.shape
batch_size = numpy.prod(batch_shape)
a = a.reshape(batch_size, m, lda)
v = cupy.array(a.swapaxes(-2, -1), order='C', copy=True, dtype=inp_v_dtype)
w = cupy.empty((batch_size, m), inp_w_dtype).swapaxes(-2, 1)
dev_info = cupy.empty((), numpy.int32)
handle = device.Device().cusolver_handle
if with_eigen_vector:
jobz = cusolver.CUSOLVER_EIG_MODE_VECTOR
else:
jobz = cusolver.CUSOLVER_EIG_MODE_NOVECTOR
if UPLO == 'L':
uplo = cublas.CUBLAS_FILL_MODE_LOWER
else: # UPLO == 'U'
uplo = cublas.CUBLAS_FILL_MODE_UPPER
if dtype == 'f':
buffer_size = cusolver.ssyevjBatched_bufferSize
syevjBatched = cusolver.ssyevjBatched
elif dtype == 'd':
buffer_size = cusolver.dsyevjBatched_bufferSize
syevjBatched = cusolver.dsyevjBatched
elif dtype == 'F':
buffer_size = cusolver.cheevjBatched_bufferSize
syevjBatched = cusolver.cheevjBatched
elif dtype == 'D':
buffer_size = cusolver.zheevjBatched_bufferSize
syevjBatched = cusolver.zheevjBatched
else:
raise RuntimeError('Only float and double and cuComplex and ' +
'cuDoubleComplex are supported')
params = cusolver.createSyevjInfo()
work_size = buffer_size(handle, jobz, uplo, m, v.data.ptr, lda, w.data.ptr,
params, batch_size)
work = cupy.empty(work_size, inp_v_dtype)
syevjBatched(handle, jobz, uplo, m, v.data.ptr, lda, w.data.ptr,
work.data.ptr, work_size, dev_info.data.ptr, params,
batch_size)
cupy.linalg.util._check_cusolver_dev_info_if_synchronization_allowed(
syevjBatched, dev_info)
cusolver.destroySyevjInfo(params)
w = w.astype(ret_w_dtype, copy=False)
w = w.swapaxes(-2, -1).reshape(*batch_shape, m)
if not with_eigen_vector:
return w
v = v.astype(ret_v_dtype, copy=False)
v = v.swapaxes(-2, -1).reshape(*batch_shape, m, m)
return w, v
def syevj(a, UPLO='L', with_eigen_vector=True):
"""Eigenvalue decomposition of symmetric matrix using cusolverDn<t>syevj().
Computes eigenvalues ``w`` and (optionally) eigenvectors ``v`` of a complex
Hermitian or a real symmetric matrix.
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``.
with_eigen_vector (bool): Indicates whether or not eigenvectors
are computed.
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]``.
"""
if not check_availability('syevj'):
raise RuntimeError('syevj is not available.')
if UPLO not in ('L', 'U'):
raise ValueError('UPLO argument must be \'L\' or \'U\'')
if a.ndim > 2:
return _syevj_batched(a, UPLO, with_eigen_vector)
assert a.ndim == 2
if a.dtype == 'f' or a.dtype == 'e':
dtype = 'f'
inp_w_dtype = 'f'
inp_v_dtype = 'f'
ret_w_dtype = a.dtype
ret_v_dtype = a.dtype
elif a.dtype == 'd':
dtype = 'd'
inp_w_dtype = 'd'
inp_v_dtype = 'd'
ret_w_dtype = 'd'
ret_v_dtype = 'd'
elif a.dtype == 'F':
dtype = 'F'
inp_w_dtype = 'f'
inp_v_dtype = 'F'
ret_w_dtype = 'f'
ret_v_dtype = 'F'
elif a.dtype == 'D':
dtype = 'D'
inp_w_dtype = 'd'
inp_v_dtype = 'D'
ret_w_dtype = 'd'
ret_v_dtype = 'D'
else:
# NumPy uses float64 when an input is not floating point number.
dtype = 'd'
inp_w_dtype = 'd'
inp_v_dtype = 'd'
ret_w_dtype = 'd'
ret_v_dtype = 'd'
# Note that cuSolver assumes fortran array
v = a.astype(inp_v_dtype, order='F', copy=True)
m, lda = a.shape
w = cupy.empty(m, inp_w_dtype)
dev_info = cupy.empty((), numpy.int32)
handle = device.Device().cusolver_handle
if with_eigen_vector:
jobz = cusolver.CUSOLVER_EIG_MODE_VECTOR
else:
jobz = cusolver.CUSOLVER_EIG_MODE_NOVECTOR
if UPLO == 'L':
uplo = cublas.CUBLAS_FILL_MODE_LOWER
else: # UPLO == 'U'
uplo = cublas.CUBLAS_FILL_MODE_UPPER
if dtype == 'f':
buffer_size = cusolver.ssyevj_bufferSize
syevj = cusolver.ssyevj
elif dtype == 'd':
buffer_size = cusolver.dsyevj_bufferSize
syevj = cusolver.dsyevj
elif dtype == 'F':
buffer_size = cusolver.cheevj_bufferSize
syevj = cusolver.cheevj
elif dtype == 'D':
buffer_size = cusolver.zheevj_bufferSize
syevj = cusolver.zheevj
else:
raise RuntimeError('Only float and double and cuComplex and ' +
'cuDoubleComplex are supported')
params = cusolver.createSyevjInfo()
work_size = buffer_size(handle, jobz, uplo, m, v.data.ptr, lda, w.data.ptr,
params)
work = cupy.empty(work_size, inp_v_dtype)
syevj(handle, jobz, uplo, m, v.data.ptr, lda, w.data.ptr, work.data.ptr,
work_size, dev_info.data.ptr, params)
cupy.linalg.util._check_cusolver_dev_info_if_synchronization_allowed(
syevj, dev_info)
cusolver.destroySyevjInfo(params)
w = w.astype(ret_w_dtype, copy=False)
if not with_eigen_vector:
return w
v = v.astype(ret_v_dtype, copy=False)
return w, v
