o װiP@sddlZddlZddlmZddlmZddlmZddlm Z ddZ dd Z d d Z d d Z dddZddZdddZdS)N)runtime)internal)device)_utilc Cs,ddlm}ddlm}|j}|d}|j|ddd||}|jd}tj||ftj d}tj|ftj d}|d ||k} | rt } |} || |j } |j j} | | |}tj| || tjd}|tjkrk|j}n|tjkrt|j}n|tjkr}|j}n |tjkr|j}nJ|| ||j j| |j j|j j|ngt } |tjkr|j}|j}n&|tjkr|j}|j}n|tjkr|j}|j}n|tjkr|j }|j!}nJt"|D],}||j j}|| ||||}tj||d}|| |||||j j||j j||j jq||||d d||d dfS) aCompute pivoted LU decomposition. Decompose a given batch of square matrices. Inputs and outputs are transposed. Args: a_t (cupy.ndarray): The input matrix with dimension ``(..., N, N)``. The dimension condition is not checked. dtype (numpy.dtype): float32, float64, complex64, or complex128. Returns: tuple: lu_t (cupy.ndarray): ``L`` without its unit diagonal and ``U`` with dimension ``(..., N, N)``. piv (cupy.ndarray): 1-origin pivot indices with dimension ``(..., N)``. dev_info (cupy.ndarray): ``getrf`` info with dimension ``(...)``. .. seealso:: :func:`scipy.linalg.lu_factor` rcublascusolverC)orderdtypeiFN)#cupy_backends.cuda.libsrr shapeastypereshapecupyemptynumpyint32rget_cublas_handleitemsizedataptrarangeuintpfloat32 sgetrfBatchedfloat64 dgetrfBatched complex64 cgetrfBatched complex128 zgetrfBatchedget_cusolver_handlesgetrf_bufferSizesgetrfdgetrf_bufferSizedgetrfcgetrf_bufferSizecgetrfzgetrf_bufferSizezgetrfrange)a_trrr orig_shapen batch_sizeipivdev_info use_batchedhandleldastepstartstopa_array getrfBatchedgetrf_bufferSizegetrfia_ptr buffersize workspacerDN/home/ubuntu/.local/lib/python3.10/site-packages/cupy/linalg/_decomposition.py _lu_factor sn              rFcCs6ddlm}ddlm}ddlm}|dstdt|\}}|jdkr-t |j |S|dkr5|j }n|dkr=|j }n |d krE|j}n|j}|j|d d d }t j|}|j d } |jd|jj} t} t|j dd} t j | tjd} || |j| |jj| | jj| t jj || t !|j|ddS)aBatched Cholesky decomposition. Decompose a given array of two-dimensional square matrices into ``L * L.T``, where ``L`` is a lower-triangular matrix and ``.T`` is a conjugate transpose operator. Args: a (cupy.ndarray): The input array of matrices with dimension ``(..., N, N)`` Returns: cupy.ndarray: The lower-triangular matrix. rrr)check_availability potrfBatchedzpotrfBatched is not availablefdFr Tr copyr r NrFrM)"rrr cupyx.cusolverrG RuntimeErrorrlinalg_common_typesizerrr spotrfBatched dpotrfBatched cpotrfBatched zpotrfBatchedr_core _mat_ptrsstridesrrrr&rprodrrCUBLAS_FILL_MODE_UPPERrrlinalg3_check_cusolver_dev_info_if_synchronization_allowedtril)arr rGr out_dtyperHxxpr2ldxr7r3r5rDrDrE_potrf_batchedjs<      rdc Cs`ddlm}ddlm}t|t|t||jdkr$t|St |\}}|j dkr7t |j |S|j|ddd}t|}t}t j dtjd }|d krZ|j} |j} n|d kre|j} |j} n|d krp|j} |j} n|j} |j} | ||j||jj|} t j | |d } | ||j||jj|| jj| |jjt jj | |tj!|dd |j|ddS)a6Cholesky decomposition. Decompose a given two-dimensional square matrix into ``L * L.H``, where ``L`` is a lower-triangular matrix and ``.H`` is a conjugate transpose operator. Args: a (cupy.ndarray): Hermitian (symmetric if all elements are real), positive-definite input matrix with dimension ``(..., M, M)``. Returns: cupy.ndarray: The lower-triangular matrix of shape ``(..., M, M)``. .. 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` rrrr TrLrrIrJrK)kFrN)"rrr r_assert_cupy_array_assert_stacked_2d_assert_stacked_squarendimrdrQrRrrrrlenrr&rrspotrfspotrf_bufferSizedpotrfdpotrf_bufferSizecpotrfcpotrf_bufferSizezpotrfzpotrf_bufferSizer[rrr\r]_tril) r_rr rr`rar2r7r5potrfpotrf_bufferSizerBrCrDrDrEcholeskysL       rxcCsddlm}|jdd}t|}|jdd\}}t||}|dks(|dkrt|\}} |dkrGt |||f| t |||f| fS|dkr^t ||| } | t |||f| fS|dkrlt |||f| S|dkrt |||f| t ||f| fS|j dg|jddR}|||} |dkr| || jddS| \} } | || jdd} |dkrdnd} | || j| d} | | fS) Nr)_geqrf_orgqr_batchedr reducedcompleterrawr ) rOryrrrZminrrQrrstacked_identityr)r_modery batch_shaper3mr2rgrr`qoutr|idxrDrDrE _qr_batcheds:    rrzc Csddlm}t||dvr%|dvrd|}t|d|}t||jdkr/t||St|\}}|j \}}t ||}|dkr|dkrXt |df|t d|f|fS|d krjt ||t ||f|fS|d krvt d|f|St ||f|t d |fS|j|d d d} t} t j dtjd} |dkr|j} |j} n+|dkr|j} |j} n |dkr|j} |j} n|dkr|j} |j} n d|j}t|| | ||| jj|}t j ||d}t j ||d}| | ||| jj||jj|jj|| jj t jj | | |d kr| ddd|f}t!|j|ddS|dkr2| j|dd|j|ddfS|d krG||krG|}t ||f|}n |}t ||f|}| |d|<|dkrc|j"}|j#}n#|dkro|j$}|j%}n|dkr{|j&}|j'}n |dkr|j(}|j)}|| ||||jj||jj}t j ||d}|| ||||jj||jj|jj|| jj t jj || |d|}| ddd|f}|j|ddt!|j|ddfS)aQR decomposition. Decompose a given two-dimensional matrix into ``Q * R``, where ``Q`` is an orthonormal and ``R`` is an upper-triangular matrix. Args: a (cupy.ndarray): The input matrix. mode (str): The mode of decomposition. Currently 'reduced', 'complete', 'r', and 'raw' modes are supported. The default mode is 'reduced', in which matrix ``A = (..., M, N)`` is decomposed into ``Q``, ``R`` with dimensions ``(..., M, K)``, ``(..., K, N)``, where ``K = min(M, N)``. Returns: cupy.ndarray, or tuple of ndarray: Although the type of returned object depends on the mode, it returns a tuple of ``(Q, R)`` by default. For details, please see the document of :func:`numpy.linalg.qr`. .. 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.qr` rr)rzr{r|r})rIfulleeconomicz)The deprecated mode '{}' is not supportedzUnrecognized mode '{}'rerzr{r|rr TrLrfrrIrJrKDzDdtype must be float32, float64, complex64 or complex128 (actual: {})NFrNr})*rr rrhformat ValueErrorrkrrQrr~rridentity transposerrr&rrsgeqrf_bufferSizesgeqrfdgeqrf_bufferSizedgeqrfcgeqrf_bufferSizecgeqrfzgeqrf_bufferSizezgeqrfrrrr\r]_triusorgqr_bufferSizesorgqrdorgqr_bufferSizedorgqrcungqr_bufferSizecungqrzungqr_bufferSizezungqr)r_rr msgrr`rr2rgrar7r5geqrf_bufferSizegeqrfrBrCtaur|mcrorgqr_bufferSizeorgqrrDrDrEqrs                   rcCsddlm}m}|jdd}t|}|jdd\}}t|\} } | j } |dkrst ||} t || f| } |rq|rVt j |||f| d}t j |||f| d}nt j ||| f| d}t j || |f| d}|| |fS| S|dks{|dkrt |d| } |r|rt ||| }t ||| }nt j ||df| d}t j |d|f| d}|| |fS| S|jdg|jddR}tjs|dkr|dkr|j| dd d }||||d }n ||| j||d }|r:|\}} }|j| d d }|jg||jddR}| j| d d } | jg|| jddR} |j| d d }|jg||jddR}|| |ddfS|} | j| d d } | jg|| jddR} | S) Nr)_gesvdj_batched_gesvd_batchedr rrr r FrLrN)rOrrrrrZrrQcharlowerr~rrrrris_hiprswapaxesconj)r_ full_matrices compute_uvrrrr3r2rruv_dtypes_dtypergsuvtrvrDrDrE _svd_batchedsZ      rTcCsddlm}t||jdkrt|||St|\}}|j}|j}|j \}} | dks4|dkrft d|} |rd|rMt j ||d} t j | |d} nt j |df|d} t j d| f|d} | | | fS| S| |kru|j |ddd} d }n|j \} }|j |ddd} d}|}|r|rt j | | f|d} | d d d |f} td }td }n| } t j ||f|d} td }td }| jj| jj}}n d\}}td}td}t j ||d} t}t j dtjd}|dkr|j}|j}n|dkr|j}|j}n|dkr|j}|j}n|j}|j}||| |}t j ||d}tjs d}nt j t | |d|d}|jj}||||| || jj| | jj|| |||jj|||jjt j!j"||| j |d d} |r|| j |d d} | j |d d} |rw| | | fS| | | fS| S)aSingular Value Decomposition. Factorizes the matrix ``a`` as ``u * np.diag(s) * v``, where ``u`` and ``v`` are unitary and ``s`` is an one-dimensional array of ``a``'s singular values. Args: a (cupy.ndarray): The input matrix with dimension ``(..., M, N)``. full_matrices (bool): If True, it returns u and v with dimensions ``(..., M, M)`` and ``(..., N, N)``. Otherwise, the dimensions of u and v are ``(..., M, K)`` and ``(..., K, N)``, respectively, where ``K = min(M, N)``. compute_uv (bool): If ``False``, it only returns singular values. Returns: tuple of :class:`cupy.ndarray`: A tuple of ``(u, s, v)`` such that ``a = u * np.diag(s) * v``. .. 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`. .. note:: On CUDA, when ``a.ndim > 2`` and the matrix dimensions <= 32, a fast code path based on Jacobian method (``gesvdj``) is taken. Otherwise, a QR method (``gesvd``) is used. On ROCm, there is no such a fast code path that switches the underlying algorithm. .. seealso:: :func:`numpy.linalg.svd` rrrerrr TrLFNAOS)rrNrfrIrJrKrN)#rr rrhrkrrQrrrrreyerrordrrrr&rrsgesvdsgesvd_bufferSizedgesvddgesvd_bufferSizecgesvdcgesvd_bufferSizezgesvdzgesvd_bufferSizerrr~r\r])r_rrr rr real_dtyperr2rrrrra trans_flagrgjob_ujob_vtu_ptrvt_ptrr7r5gesvdgesvd_bufferSizerBrC rwork_ptrrworkrDrDrEsvds $             r)rz)TT)rrcupy_backends.cuda.apir cupy._corer cupy.cudar cupy.linalgrrFrdrxrrrrrDrDrDrEs    `3D &B