U
    ¿|e€g ã                   @   sX  d Z ddlZddlZddlmZ ddlZddlZddlm	Z	m
Z
mZmZ ddlmZ ddlmZmZmZ ddlmZmZ ddlmZmZmZ d	d
lmZmZmZ ddlmZ  e !d¡Z"e" #¡ Z$e$Z%e !d¡Z&e& #¡ Z'ej(Z(e( #¡ Z)ej*Z+ej*Z,ej-dej.dej/dej0diZ1d¼dd„Z2dd„ Z3dd„ Z4dd„ Z5G dd„ dƒZ6G dd„ dƒZ7ej8dd „ ƒZ9d!d"„ Z:d#d$„ Z;d%d&„ Z<d'd(„ Z=d)d*„ Z>d+d,„ Z?d-d.„ Z@d/d0„ ZAd1d2„ ZBd½d4d5„ZCeejDƒd6d7„ ƒZEeejFƒd8d9„ ƒZGd:d;„ ZHeejIƒd<d=„ ƒZId>d?„ ZJd@dA„ ZKdBdC„ ZLdDdE„ ZMeejDƒdFdG„ ƒZNe OdHe P¡ ¡ZQedIdJ„ ƒZRd¾dLdM„ZSdNdO„ ZTdPdQ„ ZUeeUƒdRdS„ ƒZVedTdU„ ƒZWedVdW„ ƒZXeejYjZƒdXdY„ ƒZ[edZd[„ ƒZ\d¿d\d]„Z]eejYj^ƒd^d_„ ƒZ_eejYj`ƒd`da„ ƒZaeejYjbƒdbdc„ ƒZceejYjdƒddde„ ƒZeeejYjfƒdfdg„ ƒZgeejYjhƒdÀdhdi„ƒZieejYjjƒdjdk„ ƒZkdldm„ Zleelƒdndo„ ƒZmdpdq„ Zneenƒdrds„ ƒZodtdu„ Zpeepƒdvdw„ ƒZqdxdy„ Zreerƒdzd{„ ƒZsd|d}„ Zteetƒd~d„ ƒZud€d�„ Zveevƒd‚dƒ„ ƒZweejYjxƒdÁd…d†„ƒZyd‡dˆ„ Zzeezƒd‰dŠ„ ƒZ{eejYj|ƒd‹dŒ„ ƒZ}eejYj~ƒdÂdŽd�„ƒZd�d‘„ Z€eejYj�ƒd’d“„ ƒZ‚eejYjƒƒd”d•„ ƒZ„d–d—„ Z…ee…ƒd˜d™„ ƒZ†dšd›„ Z‡ee‡ƒdœd�„ ƒZˆdždŸ„ Z‰eejYjŠƒdÃd d¡„ƒZ‹eejYjŒƒdÄd¢d£„ƒZ�ed¤d¥„ ƒZŽeejYj�ƒdÅd¦d§„ƒZ�eejYj‘ƒd¨d©„ ƒZ’eej“ƒdÆdªd«„ƒZ”dÇd¬d­„Z•ed®d¯„ ƒZ–ed°d±„ ƒZ—d²d³„ Z˜eej™ƒdÈd´dµ„ƒZšd¶d·„ Z›d¸d¹„ Zœeej�ƒdºd»„ ƒZždS )Éz.
Implementation of linear algebra operations.
é    N)Úir)Úlower_builtinÚimpl_ret_borrowedÚimpl_ret_new_refÚimpl_ret_untracked)Ú	signature)Ú	intrinsicÚoverloadÚregister_jitable)ÚtypesÚcgutils)ÚTypingErrorÚNumbaTypeErrorÚNumbaPerformanceWarningé   )Ú
make_arrayÚ_empty_nd_implÚ
array_copy)Únumpy_supporté   é    ÚsÚdÚcÚzú<BLAS function>c                 C   s$   t  | ¡}|d kr td|f ƒ‚|S )Nzunsupported dtype for %s())Ú_blas_kindsÚgetÚ	TypeError)ÚdtypeÚ	func_nameÚkind© r"   úL/var/www/website-v5/atlas_env/lib/python3.8/site-packages/numba/np/linalg.pyÚget_blas_kind0   s    
r$   c                  C   s.   zdd l } W n tk
r(   tdƒ‚Y nX d S ©Nr   z*scipy 0.16+ is required for linear algebra)Zscipy.linalg.cython_blasÚImportError©Úscipyr"   r"   r#   Úensure_blas7   s    r)   c                  C   s.   zdd l } W n tk
r(   tdƒ‚Y nX d S r%   )Úscipy.linalg.cython_lapackr&   r'   r"   r"   r#   Úensure_lapack>   s    r+   c                 C   s   |   |||¡}t ||¡S ©N)Úget_constant_genericr   Úalloca_once_value)ÚcontextÚbuilderÚtyÚvalÚconstr"   r"   r#   Úmake_constant_slotE   s    r4   c                   @   s0   e Zd ZdZdd„ Zedd„ ƒZedd„ ƒZdS )	Ú_BLASzM
    Functions to return type signatures for wrapped
    BLAS functions.
    c                 C   s
   t ƒ  d S r,   )r)   ©Úselfr"   r"   r#   Ú__init__P   s    z_BLAS.__init__c              	   C   s<   t |d|ƒ}t tjtjt |¡tjt |¡¡}t d|¡S )NÚunderlying_floatÚnumba_xxnrm2©Úgetattrr   ÚintcÚcharÚintpÚCPointerÚExternalFunction©Úclsr   ÚrtypeÚsigr"   r"   r#   r:   S   s    üz_BLAS.numba_xxnrm2c                 C   s`   t  t jt jt jt jt jt jt  |¡t  |¡t jt  |¡t jt  |¡t  |¡t j¡}t  d|¡S )NÚnumba_xxgemm©r   r=   r>   r?   r@   rA   ©rC   r   rE   r"   r"   r#   rF   ^   s"    òz_BLAS.numba_xxgemmN)Ú__name__Ú
__module__Ú__qualname__Ú__doc__r8   Úclassmethodr:   rF   r"   r"   r"   r#   r5   J   s   

r5   c                   @   sœ   e Zd ZdZdd„ Zedd„ ƒZedd„ ƒZedd	„ ƒZed
d„ ƒZ	edd„ ƒZ
edd„ ƒZedd„ ƒZedd„ ƒZedd„ ƒZedd„ ƒZedd„ ƒZdS )Ú_LAPACKzO
    Functions to return type signatures for wrapped
    LAPACK functions.
    c                 C   s
   t ƒ  d S r,   )r+   r6   r"   r"   r#   r8   y   s    z_LAPACK.__init__c              
   C   s4   t  t jt jt jt  |¡t jt  t¡¡}t  d|¡S )NÚnumba_xxgetrf©r   r=   r>   r?   r@   ÚF_INT_nbtyperA   rH   r"   r"   r#   rO   |   s    ûz_LAPACK.numba_xxgetrfc              	   C   s0   t  t jt jt  |¡t jt  t¡¡}t  d|¡S )NÚnumba_ez_xxgetrirP   rH   r"   r"   r#   rR   ‡   s    üz_LAPACK.numba_ez_xxgetric                 C   sX   t  t jt jt jt jt  |¡t jt  |¡t  |¡t  |¡t jt  |¡t j¡}t  d|¡S )NÚnumba_ez_rgeevrG   rH   r"   r"   r#   rS   ‘   s    õz_LAPACK.numba_ez_rgeevc                 C   sP   t  t jt jt jt jt  |¡t jt  |¡t  |¡t jt  |¡t j¡}t  d|¡S )NÚnumba_ez_cgeevrG   rH   r"   r"   r#   rT   ¢   s    öz_LAPACK.numba_ez_cgeevc                 C   sD   t |d|ƒ}t tjtjtjtjt |¡tjt |¡¡}t d|¡S )Nr9   Únumba_ez_xxxevdr;   )rC   r   ZwtyperE   r"   r"   r#   rU   ²   s    úz_LAPACK.numba_ez_xxxevdc                 C   s,   t  t jt jt jt  |¡t j¡}t  d|¡S )NÚnumba_xxpotrfrG   rH   r"   r"   r#   rV   ¿   s    üz_LAPACK.numba_xxpotrfc                 C   s\   t |d|ƒ}t tjtjtjtjt |¡tjt |¡t |¡tjt |¡tj¡}t d|¡S )Nr9   Únumba_ez_gesddr;   )rC   r   ÚstyperE   r"   r"   r#   rW   É   s    õz_LAPACK.numba_ez_gesddc              
   C   s4   t  t jt jt jt  |¡t jt  |¡¡}t  d|¡S )NÚnumba_ez_geqrfrG   rH   r"   r"   r#   rY   Ü   s    úz_LAPACK.numba_ez_geqrfc                 C   s8   t  t jt jt jt jt  |¡t jt  |¡¡}t  d|¡S )NÚnumba_ez_xxgqrrG   rH   r"   r"   r#   rZ   è   s    ù	z_LAPACK.numba_ez_xxgqrc                 C   s^   t |d|ƒ}t tjtjtjtjt |¡tjt |¡tjt |¡tjt tj¡¡}t d|¡S )Nr9   Únumba_ez_gelsd)r<   r   r=   r>   r?   r@   Úfloat64rA   rB   r"   r"   r#   r[   õ   s    
õz_LAPACK.numba_ez_gelsdc                 C   s@   t  t jt jt jt  |¡t jt  t¡t  |¡t j¡}t  d|¡S )NÚnumba_xgesvrP   rH   r"   r"   r#   r]     s    ø
z_LAPACK.numba_xgesvN)rI   rJ   rK   rL   r8   rM   rO   rR   rS   rT   rU   rV   rW   rY   rZ   r[   r]   r"   r"   r"   r#   rN   s   s0   


	



	



rN   c                 c   sÈ   g }g }g }t |j|ƒD ]r\}}t|tjƒr6|jdkrB|| }	}
n4|jdd�}	t|	|ƒ}t| |||fƒ}
| 	|	|
f¡ | 	|	¡ | 	|
¡ qt|j
f|žŽ t|ƒfV  |D ]\}}| j |||¡ qªdS )zƒ
    Ensure that all array arguments are contiguous, if necessary by
    copying them.
    A new (sig, args) tuple is yielded.
    ÚCFÚC©ÚlayoutN)ÚzipÚargsÚ
isinstancer   ÚArrayra   Úcopyr   r   ÚappendÚreturn_typeÚtupleÚnrtÚdecref)r/   r0   rE   rc   ZnewtysÚnewargsZcopiesr1   r2   ÚnewtyÚnewvalZcopysigr"   r"   r#   Úmake_contiguous  s    

ro   c                    s0   d‰ ‡ fdd„}|   ||ttjtjƒ|f¡ dS )z.
    Check whether *n* fits in a C `int`.
    iÿÿÿc                    s   | ˆ krt dƒ‚d S )Nz$array size too large to fit in C int)ÚOverflowError)Ún©Z_maxintr"   r#   Úimpl5  s    zcheck_c_int.<locals>.implN)Úcompile_internalr   r   Únoner?   )r/   r0   rq   rs   r"   rr   r#   Úcheck_c_int/  s     ÿrv   c              	   C   sB   |j t ||¡dd��" |  |¡}| ¡  | d¡ W 5 Q R X dS )z[
    Check the integer error return from one of the BLAS wrappers in
    _helperlib.c.
    F©Úlikelyz#BLAS wrapper returned with an errorN©Úif_thenr   Úis_not_nullÚget_python_apiÚ
gil_ensureÚfatal_error©r/   r0   ÚresÚpyapir"   r"   r#   Úcheck_blas_return=  s    
r‚   c              	   C   sB   |j t ||¡dd��" |  |¡}| ¡  | d¡ W 5 Q R X dS )z]
    Check the integer error return from one of the LAPACK wrappers in
    _helperlib.c.
    Frw   z%LAPACK wrapper returned with an errorNry   r   r"   r"   r#   Úcheck_lapack_returnI  s    
rƒ   c                 C   s–   t  t  d¡ttttttg¡}t |j|d¡}	t	|ƒ}
t  
tt|
ƒ¡}t  
tt|ƒ¡}| |	|||| |t¡| |t¡| |t¡f¡}t| ||ƒ dS )zQ
    Call the BLAS vector * vector product function for the given arguments.
    r   Znumba_xxdotN)r   ÚFunctionTypeÚIntTypeÚll_charÚintp_tÚ	ll_void_pr   Úget_or_insert_functionÚmoduler$   ÚConstantÚordÚintÚcallÚbitcastr‚   )r/   r0   Ú	conjugater   rq   Úa_dataÚb_dataÚout_dataÚfntyÚfnr!   Úkind_valr€   r"   r"   r#   Ú
call_xxdotU  s"      ÿÿ


ýr—   c                 C   s  t  t  d¡ttttttttttg
¡}t |j|d¡}	|j	}
t
| ||
dƒ}t
| ||
dƒ}|jdkrt|\}}|d }n|\}}|d }t|
ƒ}t  tt|ƒ¡}t  t|r®tdƒntd	ƒ¡}| |	||||| |t¡| |t¡|| |t¡| |t¡| |t¡f
¡}t| ||ƒ d
S )zQ
    Call the BLAS matrix * vector product function for the given arguments.
    r   Znumba_xxgemvç      ð?ç        ÚFr   r   Útrq   N)r   r„   r…   r†   r‡   rˆ   r   r‰   rŠ   r   r4   ra   r$   r‹   rŒ   rŽ   r�   r‚   )r/   r0   Údo_transZm_typeÚm_shapesÚm_dataÚv_datar“   r”   r•   r   ÚalphaÚbetaÚmrq   Úldar!   r–   Útransr€   r"   r"   r#   Úcall_xxgemvk  sB         ýÿ



 


ûr¥   c           !         s8  t  t  d¡ttttttttttttttg¡}t ˆ j|d¡}|\}}|\}}|j	}t
| ˆ |dƒ}t
| ˆ |dƒ}t  ttdƒ¡‰t  ttdƒ¡‰‡ ‡‡‡fdd„}||||ƒ\}}}||||ƒ\}}}|ˆ|	|
ƒ\}}}t|ƒ}t  tt|ƒ¡}ˆ  |||||||ˆ  |t¡||||ˆ  |t¡||f¡} t| ˆ | ƒ d	S )
zQ
    Call the BLAS matrix * matrix product function for the given arguments.
    r   rF   r˜   r™   r›   rq   c                    s8   | j ˆj krˆnˆ| j dkr$|d n|d ˆ  |t¡fS )Nr_   r   r   )ra   r�   rˆ   )r1   ÚshapesÚdata©r0   ZnotransÚout_typer¤   r"   r#   Úget_array_paramª  s    
úz$call_xxgemm.<locals>.get_array_paramN)r   r„   r…   r†   r‡   rˆ   r   r‰   rŠ   r   r4   r‹   rŒ   r$   rŽ   r�   r‚   )!r/   r0   Zx_typeÚx_shapesÚx_dataÚy_typeÚy_shapesÚy_datar©   Ú
out_shapesr“   r”   r•   r¢   ÚkÚ_krq   r   r    r¡   rª   Ztransar£   Údata_aZtransbÚldbÚdata_bÚ_ZldcZdata_cr!   r–   r€   r"   r¨   r#   Úcall_xxgemm�  sT            ûÿ

    
 ýr·   c                 C   s(   dd„ }|   ||||¡}t| ||j|ƒS )z 
    np.dot(matrix, matrix)
    c                 S   sN   | j \}}|j \}}|dkr.t ||f| j¡S t ||f| j¡}t | ||¡S ©Nr   ©ÚshapeÚnpÚzerosr   ÚemptyÚdot)ÚaÚbr¢   r±   r²   rq   Úoutr"   r"   r#   Údot_implÆ  s    

zdot_2_mm.<locals>.dot_impl©rt   r   rh   ©r/   r0   rE   rc   rÂ   r€   r"   r"   r#   Údot_2_mmÂ  s    rÅ   c                 C   s(   dd„ }|   ||||¡}t| ||j|ƒS )z 
    np.dot(vector, matrix)
    c                 S   sH   | j \}|j \}}|dkr*t |f| j¡S t |f| j¡}t | ||¡S r¸   r¹   )r¿   rÀ   r¢   Ú_mrq   rÁ   r"   r"   r#   rÂ   Ö  s    
zdot_2_vm.<locals>.dot_implrÃ   rÄ   r"   r"   r#   Údot_2_vmÒ  s    rÇ   c                 C   s(   dd„ }|   ||||¡}t| ||j|ƒS )z 
    np.dot(matrix, vector)
    c                 S   sH   | j \}}|j \}|dkr*t |f| j¡S t |f| j¡}t | ||¡S r¸   r¹   )r¿   rÀ   r¢   rq   Ú_nrÁ   r"   r"   r#   rÂ   æ  s    
zdot_2_mv.<locals>.dot_implrÃ   rÄ   r"   r"   r#   Údot_2_mvâ  s    rÉ   Fc              	   C   s°   |j \}}|j}t|ƒ| ||d ƒ}t|ƒ| ||d ƒ}	t ||j¡\}
dd„ }|  ||ttj	f|j žŽ |¡ t
| ||
ƒ t ||  |¡¡}t| ||||
|j|	j|ƒ | |¡S )z<
    np.dot(vector, vector)
    np.vdot(vector, vector)
    r   r   c                 S   s$   | j \}|j \}||kr tdƒ‚d S )Nz;incompatible array sizes for np.dot(a, b) (vector * vector)©rº   Ú
ValueError)r¿   rÀ   r¢   rq   r"   r"   r#   Ú
check_argsý  s    zdot_2_vv.<locals>.check_args)rc   rh   r   r   Úunpack_tuplerº   rt   r   r   ru   rv   Úalloca_onceÚget_value_typer—   r§   Úload)r/   r0   rE   rc   r�   ÚatyÚbtyr   r¿   rÀ   rq   rÌ   rÁ   r"   r"   r#   Údot_2_vvò  s    
 ÿrÓ   c                 C   s   t d| |ƒS )z
    np.dot(a, b)
    znp.dot()©Ú
dot_2_impl©ÚleftÚrightr"   r"   r#   Údot_2  s    rÙ   c                 C   s   t d| |ƒS )z
    a @ b
    z'@'rÔ   rÖ   r"   r"   r#   Úmatmul_2  s    rÚ   c                    sd   t |tjƒr`t |tjƒr`t‡fdd„ƒ‰ |jdks<|jdkrTt dˆ||ff t¡ ‡ fdd„S d S )Nc                    s˜   |j |j f‰ ‡ fdd„}|j|jkr0tdˆ ƒ‚ˆ dkrJt |jdd¡}n>ˆ dksZˆ dkrlt |jd	d¡}nˆ d
kr||j}ntdˆ ƒ‚t|||ƒ|fS )Nc              
      s¸   t ƒ  t| |||ƒ�š\}}ˆ dkr<t| |||ƒW  5 Q R £ S ˆ dkr^t| |||ƒW  5 Q R £ S ˆ dkr€t| |||ƒW  5 Q R £ S ˆ dkr¢t| |||ƒW  5 Q R £ S tdƒ‚W 5 Q R X d S )N©é   rÜ   ©rÜ   r   ©r   rÜ   ©r   r   Úunreachable)r)   ro   rÅ   rÉ   rÇ   rÓ   ÚAssertionError©r/   r0   rE   rc   ©Úndimsr"   r#   Ú_dot2_codegen#  s    z0dot_2_impl.<locals>._impl.<locals>._dot2_codegenz)%s arguments must all have the same dtyperÛ   rÜ   r_   rÝ   rÞ   r   rß   z*%s: inputs must have compatible dimensions)Úndimr   r   r   re   r   )Útypingcontextr×   rØ   rå   rh   )Únamerã   r#   Ú_impl  s     ÿÿzdot_2_impl.<locals>._implr^   z/%s is faster on contiguous arrays, called on %sc                    s
   ˆ | |ƒS r,   r"   rÖ   ©ré   r"   r#   Ú<lambda>F  ó    zdot_2_impl.<locals>.<lambda>©rd   r   re   r   ra   ÚwarningsÚwarnr   )rè   r×   rØ   r"   )ré   rè   r#   rÕ     s    ! ÿþrÕ   c                    s^   t | tjƒrZt |tjƒrZtdd„ ƒ‰ | jdks8|jdkrNt d| |ff t¡ ‡ fdd„S dS )z
    np.vdot(a, b)
    c                 S   sJ   dd„ }|j dks|j dkr$tdƒ‚|j|jkr8tdƒ‚t|j||ƒ|fS )Nc              
   S   sB   t ƒ  t| |||ƒ�$\}}t| |||dd�W  5 Q R £ S Q R X d S )NT)r�   )r)   ro   rÓ   râ   r"   r"   r#   ÚcodegenQ  s    z$vdot.<locals>._impl.<locals>.codegenr   z&np.vdot() only supported on 1-D arraysz0np.vdot() arguments must all have the same dtype)ræ   r   r   r   )rç   r×   rØ   rð   r"   r"   r#   ré   O  s    ÿzvdot.<locals>._implr^   ú6np.vdot() is faster on contiguous arrays, called on %sc                    s
   ˆ | |ƒS r,   r"   rÖ   rê   r"   r#   rë   e  rì   zvdot.<locals>.<lambda>Nrí   rÖ   r"   rê   r#   ÚvdotI  s    
ÿþrò   c                 C   s:   | j \}|j \}}||kr"tdƒ‚|j |fkr6tdƒ‚d S )Nz;incompatible array sizes for np.dot(a, b) (vector * matrix)zFincompatible output array size for np.dot(a, b, out) (vector * matrix)rÊ   )r¿   rÀ   rÁ   r¢   rÆ   rq   r"   r"   r#   Údot_3_vm_check_argsh  s    
ró   c                 C   s:   | j \}}|j \}||kr"tdƒ‚|j |fkr6tdƒ‚d S )Nz;incompatible array sizes for np.dot(a, b) (matrix * vector)zFincompatible output array size for np.dot(a, b, out) (matrix * vector)rÊ   )r¿   rÀ   rÁ   r¢   rÈ   rq   r"   r"   r#   Údot_3_mv_check_argss  s    
rô   c                 C   sð  |j \}}}||jkst‚|j}t|ƒ| ||d ƒ}t|ƒ| ||d ƒ}	t|ƒ| ||d ƒ}
t ||j¡}t ||	j¡}t ||
j¡}|j|jk rÈ|}|}|d }|d }|j	dk}|	j
|j
 }}t}n4|}|}|d }|d }|j	dk}|j
|	j
 }}t}|  ||ttjf|j žŽ |¡ |D ]}t| ||ƒ �q|  tjd¡}| d||¡}| d||¡}| ||¡}|j|dd��`\}}|�$ t ||
j
| |
j|
j¡d¡ W 5 Q R X |� t| |||||||
j
ƒ W 5 Q R X W 5 Q R X t| ||j|
 ¡ ƒS )	zE
    np.dot(vector, matrix, out)
    np.dot(matrix, vector, out)
    r   r   rÜ   rš   r_   ú==Frw   )rc   rh   rá   r   r   r   rÍ   rº   ræ   ra   r§   ró   rô   rt   r   r   ru   rv   Úget_constantr?   Úicmp_signedÚor_Úif_elseÚmemsetÚmulÚitemsizeÚnitemsr¥   r   Ú	_getvalue)r/   r0   rE   rc   ÚxtyÚytyÚouttyr   ÚxÚyrÁ   r«   r®   r°   Zmtyr�   Zv_shaper£   rœ   rž   rŸ   rÌ   r2   ÚzeroÚ
both_emptyZmatrix_emptyÚis_emptyr½   Únonemptyr"   r"   r#   Údot_3_vm~  s`    

 ÿ
 ÿ ÿ
ÿr  c           '      C   s  |j \}}}||jkst‚|j}t|ƒ| ||d ƒ}t|ƒ| ||d ƒ}	t|ƒ| ||d ƒ}
t ||j¡}t ||	j¡}t ||
j¡}|\}}|\}}|jdks¤t‚dd„ }|  	||t
tjf|j žŽ |¡ t| ||ƒ t| ||ƒ t| ||ƒ |j}|	j}|
j}|  tjd¡}| d||¡}| d||¡}| d||¡}| || ||¡¡}|j|dd	���š\}}|�$ t ||
j| |
j|
j¡d¡ W 5 Q R X |��V |  tjd¡}| d||¡}| d||¡}| |¡��\} }!| �x | |¡�b\}"}#|"� t| |d|||||ƒ W 5 Q R X |#�( |j|jk}$t| ||$|||||ƒ W 5 Q R X W 5 Q R X W 5 Q R X |!�~ | |¡�h\}%}&|%�( |j|jk}$t| ||$|||||ƒ W 5 Q R X |&�" t| ||||||||||ƒ W 5 Q R X W 5 Q R X W 5 Q R X W 5 Q R X W 5 Q R X W 5 Q R X t| ||j|
 ¡ ƒS )
z%
    np.dot(matrix, matrix, out)
    r   r   rÜ   r_   c                 S   s>   | j \}}|j \}}||kr$tdƒ‚|j ||fkr:tdƒ‚d S )Nz;incompatible array sizes for np.dot(a, b) (matrix * matrix)zFincompatible output array size for np.dot(a, b, out) (matrix * matrix)rÊ   )r¿   rÀ   rÁ   r¢   r±   r²   rq   r"   r"   r#   rÌ   Ë  s    

zdot_3_mm.<locals>.check_argsrõ   Frw   )rc   rh   rá   r   r   r   rÍ   rº   ra   rt   r   r   ru   rv   r§   rö   r?   r÷   rø   rù   rú   rû   rü   rý   r—   r¥   r·   r   rþ   )'r/   r0   rE   rc   rÿ   r   r  r   r  r  rÁ   r«   r®   r°   r¢   r±   r²   rq   rÌ   r¬   r¯   r“   r  r  Zx_emptyZy_emptyr  r½   r  ÚoneZis_left_vecZis_right_vecZr_vecZr_matZv_vZm_vrœ   Zv_mZm_mr"   r"   r#   Údot_3_mm·  s¦    
 ÿ
 ÿ
   ÿ    ÿ"    ÿ      ý@
ÿr
  c                    st   t | tjƒrpt |tjƒrpt |tjƒrptdd„ ƒ‰ | jdksN|jdksN|jdkrdt d| |ff t¡ ‡ fdd„S dS )z
    np.dot(a, b, out)
    c                 S   s:   dd„ }|j |j ks |j |j kr(tdƒ‚t||||ƒ|fS )Nc              
   S   s–   t ƒ  t| |||ƒ�x\}}tdd„ |jd d… D ƒƒ}|dhkrZt| |||ƒW  5 Q R £ S |ddhkr€t| |||ƒW  5 Q R £ S tdƒ‚W 5 Q R X d S )Nc                 s   s   | ]}|j V  qd S r,   ©ræ   )Ú.0r  r"   r"   r#   Ú	<genexpr>  s     z8dot_3.<locals>._impl.<locals>.codegen.<locals>.<genexpr>rÜ   r   rà   )r)   ro   Úsetrc   r
  r  rá   )r/   r0   rE   rc   rä   r"   r"   r#   rð     s    
z%dot_3.<locals>._impl.<locals>.codegenz/np.dot() arguments must all have the same dtype)r   r   r   )rç   r×   rØ   rÁ   rð   r"   r"   r#   ré     s    ÿzdot_3.<locals>._implr^   rñ   c                    s   ˆ | ||ƒS r,   r"   ©r×   rØ   rÁ   rê   r"   r#   rë   1  rì   zdot_3.<locals>.<lambda>Nrí   r  r"   rê   r#   Údot_3  s    
ÿ
ÿÿþr  Znumba_fatal_errorc                 C   s.   t  | ¡D ]}t  | ¡ ¡s
t j d¡‚q
d S )Nz$Array must not contain infs or NaNs.)r»   ÚnditerÚisfiniteÚitemÚlinalgÚLinAlgError)r¿   Úvr"   r"   r#   Ú_check_finite_matrix7  s
    ÿr  Tc                 C   s�   |rdnd}||f}t | tjƒr&| j} t | tjƒsFd| }t|dd�‚| jdksdd| }t|dd�‚t | jtjtj	fƒsŒd| }t|dd�‚d S )	Nú	np.linalgr»   z&%s.%s() only supported for array typesF©ÚhighlightingrÜ   z%%s.%s() only supported on 2-D arrays.ú3%s.%s() only supported on float and complex arrays.)
rd   r   ÚOptionalÚtypere   r   ræ   r   ÚFloatÚComplex)r¿   r    Ú	la_prefixÚprefixÚinterpÚmsgr"   r"   r#   Ú_check_linalg_matrix?  s    
ÿr$  c                 G   s>   |d j }|dd … D ]"}|j |krd|  }t|dd�‚qd S )Nr   r   zAnp.linalg.%s() only supports inputs that have homogeneous dtypes.Fr  )r   r   )r    r   Út0r›   r#  r"   r"   r#   Ú_check_homogeneous_typesT  s
    

r&  c                   C   s   d S r,   r"   r"   r"   r"   r#   Ú_copy_to_fortran_order\  s    r'  c                    s&   | j dk‰| j dk‰ ‡ ‡fdd„}|S )Nrš   ÚAc                    sJ   ˆrt  | ¡}n6ˆ r<| jj}|r0t  | j¡j}qFt  | ¡}n
t  | ¡}|S r,   )r»   rf   ÚflagsÚf_contiguousÚTÚasfortranarray)r¿   ÚacpyZflag_f©ZA_layoutZF_layoutr"   r#   rs   f  s    
z&ol_copy_to_fortran_order.<locals>.implr`   )r¿   rs   r"   r.  r#   Úol_copy_to_fortran_order`  s    

r/  c                 C   s6   | dkr2| dk rt ƒ  dst‚| dkr2tj d¡‚d S )Nr   z(Matrix is singular to machine precision.)Úfatal_error_funcrá   r»   r  r  ©Úrr"   r"   r#   Ú_inv_err_handler|  s    ÿr3  c                 C   s   | d S )zFpass a list of variables to be preserved through dead code eliminationr   r"   ©r¿   r"   r"   r#   Ú_dummy_liveness_func†  s    r5  c                    sP   t ƒ  t| dƒ tƒ  | j¡‰tƒ  | j¡‰tt| jdƒƒ‰ ‡ ‡‡fdd„}|S )NÚinvc                    s    | j d }| j d |kr(d}tj |¡‚t| ƒ t| ƒ}|dkrD|S tj|td�}ˆˆ |||j||jƒ}t	|ƒ ˆˆ ||j||jƒ}t	|ƒ t
|j|jgƒ |S )Néÿÿÿÿéþÿÿÿú.Last 2 dimensions of the array must be square.r   ©r   )rº   r»   r  r  r  r'  r½   ÚF_INT_nptypeÚctypesr3  r5  Úsize)r¿   rq   r#  r-  Úipivr2  ©r!   rO   Znumba_xxgetrir"   r#   Úinv_impl˜  s    
zinv_impl.<locals>.inv_impl)r+   r$  rN   rO   r   rR   rŒ   r$   )r¿   r@  r"   r?  r#   r@  Œ  s    
r@  c                 C   s2   | dkr.| dk rt ƒ  dst‚| dkr.tdƒ‚d S )Nr   z&Internal algorithm failed to converge.)r0  rá   rË   r1  r"   r"   r#   Ú%_handle_err_maybe_convergence_problem´  s    rA  c                 C   sf   |rdnd}||f}t | tjƒs,td| ƒ‚| jdksBtd| ƒ‚t | jtjtjfƒsbtd| ƒ‚d S )Nr  r»   z'%s.%s() only supported for array types rÜ   ú+%s.%s() only supported on 1 and 2-D arrays r  )rd   r   re   r   ræ   r   r  r  ©r¿   r    r   r!  r"  r"   r"   r#   Ú_check_linalg_1_or_2d_matrix¾  s    ÿ
ÿÿrD  c                    sR   t ƒ  t| dƒ tƒ  | j¡‰tt| jdƒƒ‰tdƒ‰ tdƒ}‡ ‡‡fdd„}|S )NÚcholeskyÚUÚLc                    s¢   | j d }| j d |kr(d}tj |¡‚|  ¡ }|dkr<|S ˆˆˆ ||j|ƒ}|dkr€|dk rltƒ  dslt‚|dkr€tj d¡‚t|ƒD ]}d|d |…|f< qˆ|S )Nr7  r8  r9  r   z Matrix is not positive definite.)	rº   r»   r  r  rf   r<  r0  rá   Úrange)r¿   rq   r#  rÁ   r2  Úcol©ZUPr!   rV   r"   r#   Úcho_implÜ  s&    
ÿzcho_impl.<locals>.cho_impl)r+   r$  rN   rV   r   rŒ   r$   )r¿   ZLOrK  r"   rJ  r#   rK  Ð  s    
rK  c                    sŒ   t ƒ  t| dƒ tƒ  | j¡‰tƒ  | j¡‰tt| jdƒƒ‰tdƒ‰ tdƒ‰‡ ‡‡‡fdd„}‡ ‡‡‡fdd„}t| jt	j
jƒr„|S |S d S )NÚeigÚNÚVc                    s  | j d }| j d |kr(d}tj |¡‚t| ƒ t| ƒ}d}|}tj|| jd�}tj|| jd�}tj||f| jd�}tj||f| jd�}	|dkrš||	jfS ˆˆˆ ˆ||j	||j	|j	|j	||	j	|ƒ}
t
|
ƒ t |¡rÜtdƒ‚t|j|j|	j|j|jgƒ ||	jfS )z7
        eig() implementation for real arrays.
        r7  r8  r9  r   r:  r   z.eig() argument must not cause a domain change.)rº   r»   r  r  r  r'  r½   r   r+  r<  rA  ÚanyrË   r5  r=  ©r¿   rq   r#  r-  ÚldvlÚldvrÚwrÚwiÚvlÚvrr2  ©ÚJOBVLÚJOBVRr!   rS   r"   r#   Úreal_eig_impl
  sD    

õ
ÿzeig_impl.<locals>.real_eig_implc           
         sØ   | j d }| j d |kr(d}tj |¡‚t| ƒ t| ƒ}d}|}tj|| jd�}tj||f| jd�}tj||f| jd�}|dkrŠ||jfS ˆˆˆ ˆ||j	||j	|j	||j	|ƒ}	t
|	ƒ t|j|j|j|jgƒ ||jfS )z:
        eig() implementation for complex arrays.
        r7  r8  r9  r   r:  r   )rº   r»   r  r  r  r'  r½   r   r+  r<  rA  r5  r=  ©
r¿   rq   r#  r-  rQ  rR  ÚwrU  rV  r2  ©rX  rY  r!   rT   r"   r#   Úcmplx_eig_implB  s8    

öz eig_impl.<locals>.cmplx_eig_impl©r+   r$  rN   rS   r   rT   rŒ   r$   rd   r   Úscalarsr  )r¿   rZ  r^  r"   ©rX  rY  r!   rT   rS   r#   Úeig_implü  s    
8(rb  c                    sŒ   t ƒ  t| dƒ tƒ  | j¡‰tƒ  | j¡‰tt| jdƒƒ‰tdƒ‰ tdƒ‰‡ ‡‡‡fdd„}‡ ‡‡‡fdd„}t| jt	j
jƒr„|S |S d S )NÚeigvalsrM  c                    sî   | j d }| j d |kr(d}tj |¡‚t| ƒ t| ƒ}d}d}tj|| jd�}|dkr\|S tj|| jd�}tjd| jd�}tjd| jd�}	ˆˆˆ ˆ||j||j|j|j||	j|ƒ}
t	|
ƒ t 
|¡rÎtdƒ‚t|j|j|	j|j|jgƒ |S )z;
        eigvals() implementation for real arrays.
        r7  r8  r9  r   r:  r   z2eigvals() argument must not cause a domain change.)rº   r»   r  r  r  r'  r½   r   r<  rA  rO  rË   r5  r=  rP  rW  r"   r#   Úreal_eigvals_impl}  sD    
õ
ÿz'eigvals_impl.<locals>.real_eigvals_implc           
         sÄ   | j d }| j d |kr(d}tj |¡‚t| ƒ t| ƒ}d}d}tj|| jd�}|dkr\|S tjd| jd�}tjd| jd�}ˆˆˆ ˆ||j||j|j||j|ƒ}	t	|	ƒ t
|j|j|j|jgƒ |S )z>
        eigvals() implementation for complex arrays.
        r7  r8  r9  r   r:  r   )rº   r»   r  r  r  r'  r½   r   r<  rA  r5  r=  r[  r]  r"   r#   Úcmplx_eigvals_impl¸  s8    
öz(eigvals_impl.<locals>.cmplx_eigvals_implr_  )r¿   rd  re  r"   ra  r#   Úeigvals_implo  s    
;)rf  c                    sp   t ƒ  t| dƒ t| jd| jƒ}t |¡‰tƒ  | j¡‰tt	| jdƒƒ‰tdƒ‰ tdƒ‰‡ ‡‡‡‡fdd„}|S )NÚeighr9   rN  rG  c                    sŽ   | j d }| j d |kr(d}tj |¡‚t| ƒ t| ƒ}tj|ˆd�}|dkrV||fS ˆˆˆ ˆ||j||jƒ}t|ƒ t	|j
|j
gƒ ||fS ©Nr7  r8  r9  r:  r   ©rº   r»   r  r  r  r'  r½   r<  rA  r5  r=  ©r¿   rq   r#  r-  r\  r2  ©ÚJOBZÚUPLOr!   rU   Zw_dtyper"   r#   Ú	eigh_impl÷  s(    
úzeigh_impl.<locals>.eigh_impl©
r+   r$  r<   r   Ú
np_supportÚas_dtyperN   rU   rŒ   r$   )r¿   Úw_typern  r"   rk  r#   rn  æ  s    

rn  c                    sp   t ƒ  t| dƒ t| jd| jƒ}t |¡‰tƒ  | j¡‰tt	| jdƒƒ‰tdƒ‰ tdƒ‰‡ ‡‡‡‡fdd„}|S )NÚeigvalshr9   rM  rG  c                    s†   | j d }| j d |kr(d}tj |¡‚t| ƒ t| ƒ}tj|ˆd�}|dkrR|S ˆˆˆ ˆ||j||jƒ}t|ƒ t	|j
|j
gƒ |S rh  ri  rj  rk  r"   r#   Úeigvalsh_impl(  s(    
úz$eigvalsh_impl.<locals>.eigvalsh_implro  )r¿   rr  rt  r"   rk  r#   rt    s    

rt  c                    sr   t ƒ  t| dƒ t| jd| jƒ}t |¡‰tƒ  | j¡‰tt	| jdƒƒ‰tdƒ‰ tdƒ‰d‡ ‡‡‡‡fdd„	}|S )	NÚsvdr9   r(  ÚSr   c                    sô   | j d }| j d }|dks$|dkr0tj d¡‚t| ƒ t| ƒ}|}t||ƒ}|r`ˆ }|}|}	nˆ}|}|}	tj||f| jd�}
tj|ˆd�}tj||	f| jd�}ˆˆ||||j	||j	|
j	||j	|	ƒ}t
|ƒ t|j|j|
j|jgƒ |
j||jfS )Nr7  r8  r   úArrays cannot be emptyr:  )rº   r»   r  r  r  r'  Úminr½   r   r<  rA  r5  r=  r+  )r¿   Úfull_matricesrq   r¢   r-  ÚlduÚminmnrl  ÚucolÚldvtÚur   Úvtr2  ©ZJOBZ_AZJOBZ_Sr!   rW   Ús_dtyper"   r#   Úsvd_implY  sD    


õzsvd_impl.<locals>.svd_impl)r   )
r+   r$  r<   r   rp  rq  rN   rW   rŒ   r$   )r¿   ry  Ús_typer‚  r"   r€  r#   r‚  H  s    

.r‚  c                    sP   t ƒ  t| dƒ tƒ  | j¡‰tƒ  | j¡‰tt| jdƒƒ‰ ‡ ‡‡fdd„}|S )NÚqrc                    sN  | j d }| j d }|dks$|dkr0tj d¡‚t| ƒ t| ƒ}|}t||ƒ}tj|| jd�}ˆˆ |||j	||j	ƒ}|dk rŠt
ƒ  dsŠt‚tj||f| jd�j}t|ƒD ]*}	t|	d ƒD ]}
||
|	f ||
|	f< q¸q¨t||ƒD ]&}	t|ƒD ]}
||
|	f ||
|	f< qêqÞˆˆ ||||j	||j	ƒ}t|ƒ t|j|jgƒ |d d …d |…f |fS )Nr7  r8  r   rw  r:  r   )rº   r»   r  r  r  r'  rx  r½   r   r<  r0  rá   r¼   r+  rH  rA  r5  r=  )r¿   rq   r¢   Úqr£   r{  ÚtauÚretr2  ÚiÚj©r!   rY   rZ   r"   r#   Úqr_impl›  sN    


úù	zqr_impl.<locals>.qr_impl)r+   r$  rN   rY   r   rZ   rŒ   r$   )r¿   r‹  r"   rŠ  r#   r‹  Š  s    
9r‹  c                 C   s   t ‚dS )z;
    Correctly copy 'b' into the 'bcpy' scratch space.
    N©ÚNotImplementedError©ÚbcpyrÀ   Únrhsr"   r"   r#   Ú_system_copy_in_bÛ  s    r‘  c                 C   s&   |j dkrdd„ }|S dd„ }|S d S )Nr   c                 S   s   || d |j d …df< d S )Nr7  r   ©rº   rŽ  r"   r"   r#   Ú	oneD_implå  s    z)_system_copy_in_b_impl.<locals>.oneD_implc                 S   s   || d |j d …d |…f< d S )Nr8  r’  rŽ  r"   r"   r#   Ú	twoD_implé  s    z)_system_copy_in_b_impl.<locals>.twoD_implr  )r�  rÀ   r�  r“  r”  r"   r"   r#   Ú_system_copy_in_b_implâ  s
    
r•  c                 C   s   t ‚dS )zK
    Compute the number of right hand sides in the system of equations
    NrŒ  ©rÀ   r"   r"   r#   Ú_system_compute_nrhsî  s    r—  c                 C   s&   | j dkrdd„ }|S dd„ }|S d S )Nr   c                 S   s   dS ©Nr   r"   r–  r"   r"   r#   r“  ø  s    z,_system_compute_nrhs_impl.<locals>.oneD_implc                 S   s
   | j d S )Nr7  r’  r–  r"   r"   r#   r”  ü  s    z,_system_compute_nrhs_impl.<locals>.twoD_implr  )rÀ   r“  r”  r"   r"   r#   Ú_system_compute_nrhs_implõ  s
    
r™  c                 C   s   t ‚dS )zD
    Check that AX=B style system input is dimensionally valid.
    NrŒ  ©r¿   rÀ   r"   r"   r#   Ú!_system_check_dimensionally_valid  s    r›  c                 C   s*   |j }|dkrdd„ }|S dd„ }|S d S )Nr   c                 S   s,   | j d }|j d }||kr(tj d¡‚d S )Nr8  r7  ú<Incompatible array sizes, system is not dimensionally valid.©rº   r»   r  r  ©r¿   rÀ   ÚamÚbmr"   r"   r#   r“    s    

ÿz9_system_check_dimensionally_valid_impl.<locals>.oneD_implc                 S   s,   | j d }|j d }||kr(tj d¡‚d S )Nr8  rœ  r�  rž  r"   r"   r#   r”    s    

ÿz9_system_check_dimensionally_valid_impl.<locals>.twoD_implr  ©r¿   rÀ   ræ   r“  r”  r"   r"   r#   Ú&_system_check_dimensionally_valid_impl  s    r¢  c                 C   s   t ‚dS )z:
    Check that AX=B style system input is not empty.
    NrŒ  rš  r"   r"   r#   Ú_system_check_non_empty  s    r£  c                 C   s*   |j }|dkrdd„ }|S dd„ }|S d S )Nr   c                 S   sF   | j d }| j d }|j d }|dks6|dks6|dkrBtj d¡‚d S ©Nr8  r7  r   rw  r�  )r¿   rÀ   rŸ  Úanr   r"   r"   r#   r“  (  s
    


z/_system_check_non_empty_impl.<locals>.oneD_implc                 S   sX   | j d }| j d }|j d }|j d }|dksH|dksH|dksH|dkrTtj d¡‚d S r¤  r�  )r¿   rÀ   rŸ  r¥  r   Úbnr"   r"   r#   r”  0  s    



 z/_system_check_non_empty_impl.<locals>.twoD_implr  r¡  r"   r"   r#   Ú_system_check_non_empty_impl$  s    r§  c                 C   s   t ‚dS )z:
    Compute the residual from the 'b' scratch space.
    NrŒ  )rÀ   rq   r�  r"   r"   r#   Ú_lstsq_residual:  s    r¨  c                    s�   | j }| j}t t|d|ƒ¡‰ |dkrTt|tjƒrB‡ fdd„}|S ‡ fdd„}|S n8|dks`t‚t|tjƒr|‡ fdd„}|S ‡ fd	d„}|S d S )
Nr9   r   c                    s6   t jdˆ d�}t  t  | |d …df ¡d ¡|d< |S ©N©r   r:  r   rÜ   )r»   r½   ÚsumÚabs©rÀ   rq   r�  r€   ©Ú
real_dtyper"   r#   Ú
cmplx_implI  s    $z(_lstsq_residual_impl.<locals>.cmplx_implc                    s0   t jdˆ d�}t  | |d …df d ¡|d< |S r©  )r»   r½   r«  r­  r®  r"   r#   Ú	real_implO  s    z'_lstsq_residual_impl.<locals>.real_implrÜ   c                    sD   t j|ˆ d�}t|ƒD ](}t  t  | |d …|f ¡d ¡||< q|S ©Nr:  rÜ   )r»   r½   rH  r«  r¬  ©rÀ   rq   r�  r€   r±   r®  r"   r#   r°  W  s    &c                    s>   t j|ˆ d�}t|ƒD ]"}t  | |d …|f d ¡||< q|S r²  )r»   r½   rH  r«  r³  r®  r"   r#   r±  ^  s     )	ræ   r   rp  rq  r<   rd   r   r  rá   )rÀ   rq   r�  ræ   r   r°  r±  r"   r®  r#   Ú_lstsq_residual_implA  s    r´  c                 C   s   t ‚dS )zŠ
    Extract 'x' (the lstsq solution) from the 'bcpy' scratch space.
    Note 'b' is only used to check the system input dimension...
    NrŒ  ©rÀ   r�  rq   r"   r"   r#   Ú_lstsq_solutionf  s    r¶  c                 C   s&   | j dkrdd„ }|S dd„ }|S d S )Nr   c                 S   s   |j  ¡ d |… S r,   ©r+  Úravelrµ  r"   r"   r#   r“  q  s    z'_lstsq_solution_impl.<locals>.oneD_implc                 S   s   |d |…d d …f   ¡ S r,   ©rf   rµ  r"   r"   r#   r”  u  s    z'_lstsq_solution_impl.<locals>.twoD_implr  )rÀ   r�  rq   r“  r”  r"   r"   r#   Ú_lstsq_solution_impln  s
    
rº  ç      ð¿c                    s‚   t ƒ  t| dƒ t|dƒ td| |ƒ t | j¡‰| j}t|d|ƒ}t |¡‰tƒ  	| j¡‰t
t|dƒƒ‰ d‡ ‡‡‡fdd„	}|S )NÚlstsqr9   r»  c                    s2  | j d }| j d }t|ƒ}t| ƒ t|ƒ t| |ƒ t| |ƒ t||ƒ}t||ƒ}t| ƒ}tj	||fˆd�j
}	t|	||ƒ tj	|ˆd�}
tj	dtjd�}ˆˆ ||||j||	j||
j||jƒ}t|ƒ |d }||k sÞ||krîtj	dˆd�}nt|	||ƒ}t||	|ƒ}t|j|	j|
j|jgƒ ||||
d |… fS )Nr7  r8  r:  r   r   )rº   r—  r  r£  r›  rx  Úmaxr'  r»   r½   r+  r‘  Úint32r<  rA  r¨  r¶  r5  r=  )r¿   rÀ   Úrcondrq   r¢   r�  r{  Zmaxmnr-  r�  r   Zrank_ptrr2  Úrankr€   r  ©r!   Únp_dtr[   r¯  r"   r#   Ú
lstsq_impl—  sF    





õzlstsq_impl.<locals>.lstsq_impl)r»  )r+   r$  rD  r&  rp  rq  r   r<   rN   r[   rŒ   r$   )r¿   rÀ   r¿  Únb_dtZr_typerÃ  r"   rÁ  r#   rÃ  z  s    


?rÃ  c                 C   s   t ‚dS )z„
    Extract 'x' (the solution) from the 'bcpy' scratch space.
    Note 'b' is only used to check the system input dimension...
    NrŒ  ©rÀ   r�  r"   r"   r#   Ú_solve_compute_returnÙ  s    rÆ  c                 C   s&   | j dkrdd„ }|S dd„ }|S d S )Nr   c                 S   s
   |j  ¡ S r,   r·  rÅ  r"   r"   r#   r“  ä  s    z-_solve_compute_return_impl.<locals>.oneD_implc                 S   s   |S r,   r"   rÅ  r"   r"   r#   r”  è  s    z-_solve_compute_return_impl.<locals>.twoD_implr  )rÀ   r�  r“  r”  r"   r"   r#   Ú_solve_compute_return_implá  s
    
rÇ  c                    sh   t ƒ  t| dƒ t|dƒ td| |ƒ t | j¡‰| j}tƒ  | j¡‰t	t
|dƒƒ‰ ‡ ‡‡fdd„}|S )NÚsolvec              	      s¶   | j d }t|ƒ}t| ƒ t|ƒ t| |ƒ t| ƒ}tj||fˆd�j}|dkrZt||ƒS t	|||ƒ tj|t
d�}ˆˆ |||j||j|j|ƒ}t|ƒ t|j|j|jgƒ t||ƒS )Nr7  r:  r   )rº   r—  r  r›  r'  r»   r½   r+  rÆ  r‘  r;  r<  r3  r5  r=  )r¿   rÀ   rq   r�  r-  r�  r>  r2  ©r!   rÂ  r]   r"   r#   Ú
solve_implþ  s0    


ø
zsolve_impl.<locals>.solve_impl)r+   r$  rD  r&  rp  rq  r   rN   r]   rŒ   r$   )r¿   rÀ   rÄ  rÊ  r"   rÉ  r#   rÊ  í  s    

)rÊ  çVçž¯Ò<c              
      s¼   t ƒ  t| dƒ t| jd| jƒ}t |¡‰tƒ  | j¡‰tƒ  	| j¡‰t
t| jdƒƒ‰t
dƒ‰ t
dƒ‰t
dƒ‰t | j¡}tjdg|d�‰tjdg|d�‰d‡ ‡‡‡‡‡‡‡‡f	d	d
„	}|S )NÚpinvr9   rv  r_   r™   r:  r˜   rË  c                    sâ  | j d }| j d }t| ƒ t| ƒ}|dks4|dkrH|j ¡  | j ¡jS t||ƒ}tj||f| j	d�}tj|ˆd�}tj||f| j	d�}ˆˆˆ |||j
||j
|j
||j
|ƒ}	t|	ƒ |d | }
d}t|ƒD ]$}|| |
krÌd||  ||< |}qÌ|d7 }||k�rBt|ƒD ]2}t|ƒD ]"}|||f ||  |||f< �q�qn@t|ƒD ]6}|| }t|ƒD ]}|||f | |||f< �q^�qJˆˆˆˆ|||ˆj
|j
||j
|ˆj
|j
|ƒ}	t|j|j|j|jˆjˆjgƒ |j ¡  | j ¡jS )Nr7  r8  r   r:  r˜   r   )rº   r  r'  r+  r¸  Úreshaperx  r»   r½   r   r<  rA  rH  r5  r=  )r¿   r¿  rq   r¢   r-  r{  r~  r   r  r2  Zcut_atZcut_idxr±   rˆ  r‰  Zs_local©	ZJOBZTRANSAZTRANSBr!   rW   rF   r	  r�  r  r"   r#   Ú	pinv_implD  sv    '


õ
& 	òÿzpinv_impl.<locals>.pinv_impl)rË  )r+   r$  r<   r   rp  rq  rN   rW   r5   rF   rŒ   r$   r»   Úarray)r¿   r¿  rƒ  ÚdtrÏ  r"   rÎ  r#   rÏ  *  s     

 rÏ  c                 C   s2   t | jtjƒrtdd„ ƒ}|S tdd„ ƒ}|S dS )zù
    Walks the diag of a LUP decomposed matrix
    uses that det(A) = prod(diag(lup(A)))
    and also that log(a)+log(b) = log(a*b)
    The return sign is adjusted based on the values found
    such that the log(value) stays in the real domain.
    c                 S   sV   |d }d}t | ƒD ]8}t |||f ¡}||||f |  }|t |¡ }q||fS )Ny                r™   )rH  r»   r¬  Úlog)rq   r¿   ÚsgnZcsgnÚaccr±   Zabselr"   r"   r#   Úcmplx_diag_walker×  s    z3_get_slogdet_diag_walker.<locals>.cmplx_diag_walkerc                 S   sL   d}t | ƒD ]2}|||f }|dk r0| }| }|t |¡ }q|d |fS )Nr™   )rH  r»   rÒ  )rq   r¿   rÓ  rÔ  r±   r  r"   r"   r#   Úreal_diag_walkerã  s    z2_get_slogdet_diag_walker.<locals>.real_diag_walkerN)rd   r   r   r  r
   )r¿   rÕ  rÖ  r"   r"   r#   Ú_get_slogdet_diag_walkerÎ  s    
	
r×  c                    sl   t ƒ  t| dƒ tƒ  | j¡‰tt| jdƒƒ‰t| ƒ‰|  d¡‰ t| jd| jƒdƒ‰‡ ‡‡‡‡fdd„}|S )NÚslogdetr   r9   r   c                    sÚ   | j d }| j d |kr(d}tj |¡‚|dkr8ˆ ˆfS t| ƒ t| ƒ}tj|td�}ˆˆ|||j||jƒ}|dkr€dtj	 fS t
|ƒ d}t|ƒD ]}||| |d k }q”|d@ }|dkrÂd}t|jgƒ ˆ|||ƒS )Nr7  r8  r9  r   r:  r™   r   )rº   r»   r  r  r  r'  r½   r;  r<  Úinfr3  rH  r5  r=  )r¿   rq   r#  r-  r>  r2  rÓ  r±   ©ÚONEÚZEROZdiag_walkerr!   rO   r"   r#   Úslogdet_impl  s*    
	z"slogdet_impl.<locals>.slogdet_impl)	r+   r$  rN   rO   r   rŒ   r$   r×  r<   )r¿   rÝ  r"   rÚ  r#   rÝ  ò  s    

)rÝ  c                 C   s   t ƒ  t| dƒ dd„ }|S )NÚdetc                 S   s   t j | ¡\}}|t  |¡ S r,   )r»   r  rØ  Úexp)r¿   rÓ  rØ  r"   r"   r#   Údet_impl4  s    zdet_impl.<locals>.det_impl©r+   r$  )r¿   rà  r"   r"   r#   rà  -  s    
rà  c                 C   s   t ‚dS )z)
    Compute singular values of *a*.
    NrŒ  r4  r"   r"   r#   Ú_compute_singular_values;  s    râ  c                    s‚   t ƒ  | j¡‰tt| jdƒƒ‰tdƒ‰ t| jd| jƒ}t |¡‰t | j¡}tj	d|d�‰tj	d|d�‰‡ ‡‡‡‡‡fdd„}|S )z>
    Returns a function to compute singular values of `a`
    ru  rM  r9   rß   r:  c           
         s¬   | j d }| j d }|dks$|dkr0tj d¡‚t| ƒ |}t||ƒ}d}d}t| ƒ}tj|ˆd�}ˆˆˆ |||j||jˆj|ˆj|ƒ}	t	|	ƒ t
|jˆjˆj|jgƒ |S )z+
        Computes singular values.
        r7  r8  r   rw  r   r:  )rº   r»   r  r  r  rx  r'  r½   r<  rA  r5  r=  )
r¿   rq   r¢   rz  r{  r|  r}  r-  r   r2  ©ZJOBZ_Nr!   Únp_ret_typerW   r~  r  r"   r#   Úsv_functionW  s6    


õz2_compute_singular_values_impl.<locals>.sv_function)
rN   rW   r   rŒ   r$   r<   rp  rq  r»   r½   )r¿   Únb_ret_typeÚnp_dtyperå  r"   rã  r#   Ú_compute_singular_values_implB  s    
/rè  c                 C   s   t ‚dS )z.
    Compute the L2-norm of 1D-array *a*.
    NrŒ  r4  r"   r"   r#   Ú_oneD_norm_2‰  s    ré  c                    sL   t | jd| jƒ}t |¡‰tƒ  | j¡‰tt| jdƒƒ‰ ‡ ‡‡fdd„}|S )Nr9   Únormc                    sl   t | ƒ}tjdˆd�}t| jd | j ƒ}ˆˆ || j||jƒ}|dk rTtƒ  dsTt‚t	|j
| j
gƒ |d S )Nrª  r:  r   )Úlenr»   r½   r�   Ústridesrü   r<  r0  rá   r5  r=  )r¿   rq   r‡  Újmpr2  ©r!   rä  Úxxnrm2r"   r#   rs   ™  s    ûz_oneD_norm_2_impl.<locals>.impl)r<   r   rp  rq  r5   r:   rŒ   r$   )r¿   ræ  rs   r"   rî  r#   Ú_oneD_norm_2_impl�  s    
rð  c           	         s
  t | jd| jƒ}t |¡}t | j¡}tƒ  | j¡}tt| jdƒƒ}| jdkrv|d t	j
fkrhddd„}n
ddd„}|S | jdk� rü|d t	j
fkrÜ| jdkr¨td	d
„ ƒ‰ n$| jdkrÀtdd
„ ƒ‰ ntdd
„ ƒ‰ d‡ fdd„	}nt |j¡j‰d‡fdd„	}|S d�st‚d S )Nr9   rê  r   c                 S   s   t | ƒS r,   )ré  ©r¿   rŒ   r"   r"   r#   r“  Ñ  s    z!_get_norm_impl.<locals>.oneD_implc                 S   sD  t | ƒ}|dkrdS |dkr$t| ƒS |tjkrft| d ƒ}td|ƒD ]}t| | ƒ}||krD|}qD|S |tj krªt| d ƒ}td|ƒD ]}t| | ƒ}||k rˆ|}qˆ|S |dkrÜd}t|ƒD ]}| | dkr¾|d7 }q¾|S |dk�rd}t|ƒD ]}|t| | ƒ7 }qò|S d}t|ƒD ]}|t| | ƒ| 7 }�q|d|  S d S )Nr   r™   rÜ   r   r˜   )rë  ré  r»   rÙ  r¬  rH  )r¿   rŒ   rq   r‡  r±   r2   r"   r"   r#   r“  Ô  sD    


rÜ   r_   c                 S   s   | S r,   r"   r4  r"   r"   r#   Úarray_prepare	  s    z%_get_norm_impl.<locals>.array_preparerš   c                 S   s   | j S r,   )r+  r4  r"   r"   r#   rò  	  s    c                 S   s   |   ¡ S r,   r¹  r4  r"   r"   r#   rò   	  s    c                    s(   | j }|dkrdS ˆ | ƒ}t| |¡ƒS )Nr   r™   )r=  ré  rÍ  )r¿   rŒ   rq   Za_c)rò  r"   r#   r”  &	  s
    z!_get_norm_impl.<locals>.twoD_implc           	         sª  | j d }| j d }| jdkr"dS |tjkrtd}t|ƒD ]6}d}t|ƒD ]}|t| ||f ƒ7 }qH||kr8|}q8|S |tj krÈˆ }t|ƒD ]6}d}t|ƒD ]}|t| ||f ƒ7 }qœ||k rŒ|}qŒ|S |dk�rd}t|ƒD ]6}d}t|ƒD ]}|t| ||f ƒ7 }qî||krÞ|}qÞ|S |dk�rrˆ }t|ƒD ]<}d}t|ƒD ]}|t| ||f ƒ7 }�q@||k �r0|}�q0|S |dk�rˆt| ƒd S |dk�ržt| ƒd S tdƒ‚d S )Nr7  r8  r   r™   r   rÜ   z Invalid norm order for matrices.)rº   r=  r»   rÙ  rH  r¬  râ  rË   )	r¿   rŒ   rq   r¢   Z
global_maxÚiiÚtmpÚjjZ
global_min)Úmax_valr"   r#   r”  1	  sZ    








r   )N)N)N)N)r<   r   rp  rq  r5   r:   rŒ   r$   ræ   r   ru   ra   r
   r»   Úfinfor  r½  rá   )	r¿   Zord_flagræ  rä  rç  rï  r!   r“  r”  r"   )rò  rö  r#   Ú_get_norm_impl´  s2    


9


	Erø  c                 C   s   t ƒ  t| dƒ t| |ƒS )Nrê  )r+   rD  rø  rñ  r"   r"   r#   Ú	norm_impl{	  s    
rù  c                 C   s   t ƒ  t| dƒ ddd„}|S )NÚcondc                 S   s    |dks|dks|d kr\t | ƒ}|dks0|d krFt |d |d ¡}qˆt |d |d ¡}n,tj | |¡}tj tj | ¡|¡}|| }t |¡r˜tjS |S d S )NrÜ   r8  r   r7  )râ  r»   Údivider  rê  r6  ÚisnanrÙ  )r¿   Úpr   r2  Znorm_aZ
norm_inv_ar"   r"   r#   rs   Š	  s    
zcond_impl.<locals>.impl)Nrá  )r¿   rý  rs   r"   r"   r#   Ú	cond_impl„	  s    

$rþ  c                 C   s4   d}t t| ƒƒD ]}| | |kr*|d }q q0q|S )zJ
    Gets rank from singular values with cut-off at a given tolerance
    r   r   ©rH  rë  )Úsvr›   rÀ  r±   r"   r"   r#   Ú_get_rank_from_singular_values±	  s    
r  c                    s.   t ƒ  t| dƒ dd„ ‰ ‡ fdd„}|| |ƒS )ah  
    Computes rank for matrices and vectors.
    The only issue that may arise is that because numpy uses double
    precision lapack calls whereas numba uses type specific lapack
    calls, some singular values may differ and therefore counting the
    number of them above a tolerance may lead to different counts,
    and therefore rank, in some cases.
    Úmatrix_rankc                    sX   |d t jfkrFt| jd| jƒ}t |¡}t |¡j‰ d‡ fdd„	}|S ddd„}|S d S )Nr9   c                    s@   t | ƒ}| jd }| jd }t||ƒ}|d | ˆ  }t||ƒS )Nr   r   )râ  rº   r½  r  )r¿   Útolr   r2  r   Úlr›   ©Zeps_valr"   r#   Ú_2d_tol_none_implÕ	  s    


zImatrix_rank_impl.<locals>._2d_matrix_rank_impl.<locals>._2d_tol_none_implc                 S   s   t | ƒ}t||ƒS r,   )râ  r  )r¿   r  r   r"   r"   r#   Ú_2d_tol_not_none_implß	  s    zMmatrix_rank_impl.<locals>._2d_matrix_rank_impl.<locals>._2d_tol_not_none_impl)N)N)	r   ru   r<   r   rp  rq  r»   r÷  Úeps)r¿   r  Únb_typeÚnp_typer  r  r"   r  r#   Ú_2d_matrix_rank_implÍ	  s    

z.matrix_rank_impl.<locals>._2d_matrix_rank_implc                    s:   | j }|dkrddd„}|S |dkr.ˆ | |ƒS ds6t‚d S )Nr   c                 S   s(   t t| ƒƒD ]}| | dkr dS qdS )Nr™   r   r   rÿ  )r¿   r  r±   r"   r"   r#   Ú_1d_matrix_rank_impló	  s    zMmatrix_rank_impl.<locals>._get_matrix_rank_impl.<locals>._1d_matrix_rank_implrÜ   r   )N)ræ   rá   )r¿   r  ræ   r  ©r  r"   r#   Ú_get_matrix_rank_implä	  s    

z/matrix_rank_impl.<locals>._get_matrix_rank_impl)r+   rD  )r¿   r  r  r"   r  r#   Úmatrix_rank_impl¿	  s
    

r  c                    sF   t | dƒ t | j¡‰ t|d|ƒ}t|tjƒs6tdƒ‚‡ fdd„}|S )zL
    Computes matrix power. Only integer powers are supported in numpy.
    Úmatrix_powerr   zExponent must be an integer.c           
         sJ  |dkr<t j| jˆ d�}t| jd ƒD ]}d|||f< q&|S | jd | jd  }}||krbtdƒ‚|dkrr|  ¡ S |dk ržt j | ¡ ¡ }|dkr–|S | }n|dkr®|  ¡ S | }|dk rì|d	krÎt  ||¡S |d
krêt  t  ||¡|¡S nZ|}|}|}d}	|dk�rB|d@ �r,|	�r |}d}	nt  ||¡}t  ||¡}|d? }qü|S d S )Nr   r:  r˜   r7  r8  zinput must be a square arrayr   é   rÜ   é   TF)	r»   r¼   rº   rH  rË   rf   r  r6  r¾   )
r¿   rq   r(  r±   rŸ  r¥  rÔ  rß  r‡  Úflag©rç  r"   r#   Úmatrix_power_impl
  sH    


z,matrix_power_impl.<locals>.matrix_power_impl)	r$  rp  rq  r   r<   rd   r   ÚIntegerr   )r¿   rq   Úntr  r"   r  r#   r  
  s    
?r  c                 C   s8   t | ddd� t|ttjfƒs*td| ƒ‚ddd„}|S )	z)
    Computes the trace of an array.
    ÚtraceF©r   z!integer argument expected, got %sr   c                 S   s”   | j \}}|}|dk r|| }|dkr.|| }tt||ƒdƒ}d}|dkrnt|ƒD ]}|| ||| f 7 }qRn"t|ƒD ]}|| || |f 7 }qv|S r¸   )rº   r½  rx  rH  )r¿   ÚoffsetÚrowsÚcolsr±   rq   r‡  rˆ  r"   r"   r#   Úmatrix_trace_impl]
  s    
z,matrix_trace_impl.<locals>.matrix_trace_impl)r   )r$  rd   r�   r   r  r   )r¿   r  r  r"   r"   r#   r  R
  s
    
r  c                 C   s>   |rdnd}||f}t | tjƒr:| jdks:td| dd�‚d S )Nr  r»   rÜ   rB  Fr  )rd   r   re   ræ   r   rC  r"   r"   r#   Ú_check_scalar_or_lt_2d_matq
  s    
ÿÿr  c                 C   s@   t  | ¡}t  |¡}t  | ¡  |jdf¡| ¡  d|jf¡¡S r˜  ©r»   ÚasarrayÚmultiplyr¸  rÍ  r=  ©r¿   rÀ   rÁ   ÚaaÚbbr"   r"   r#   Úouter_impl_none{
  s
    

ÿr%  c                 C   sF   t  | ¡}t  |¡}t  | ¡  |jdf¡| ¡  d|jf¡|¡ |S r˜  r  r"  r"   r"   r#   Úouter_impl_arrƒ
  s    

þr&  c                 C   s   |d t jfkrtS tS d S r,   )r   ru   r%  r&  ©r¿   rÀ   rÁ   r"   r"   r#   Ú_get_outer_impl�
  s    r(  c                    s:   t | ddd� t |ddd� t| ||ƒ‰ d‡ fdd„	}|S )NÚouterFr  c                    s   ˆ | ||ƒS r,   r"   r'  ©rs   r"   r#   Ú
outer_implœ
  s    zouter_impl.<locals>.outer_impl)N)r  r(  )r¿   rÀ   rÁ   r+  r"   r*  r#   r+  ”
  s
    r+  c                 C   sh   t | tjƒrT| jdkr(td | j¡ƒ‚qd| jdkrBtdd„ ƒ}|S tdd„ ƒ}|S ntdd„ ƒ}|S d S )N)r_   rš   z^np.linalg.kron only supports 'C' or 'F' layout input arrays. Received an input of layout '{}'.rÜ   c                 S   s    | j d }| j d }|  ||¡S )Nr7  r8  ©rº   rÍ  )r  ÚxnÚxmr"   r"   r#   Ú	nrm_shapeª
  s    

z(_kron_normaliser_impl.<locals>.nrm_shapec                 S   s   | j d }|  d|¡S )Nr7  r   r,  )r  r-  r"   r"   r#   r/  ±
  s    
c                 S   s   t  dt| ƒ¡}| |d< |S )Nrß   r   )r»   r½   r  )r  r¿   r"   r"   r#   r/  ·
  s    )rd   r   re   ra   r   Úformatræ   r
   )r  r/  r"   r"   r#   Ú_kron_normaliser_impl¢
  s    
þ



r1  c                 C   s’   t | tjƒ}t |tjƒ}|rV|rV| jdks4|jdkrDtdd„ ƒ}|S tdd„ ƒ}|S n8|rjtdd„ ƒ}|S |r~tdd„ ƒ}|S tdd„ ƒ}|S d S )NrÜ   c                 S   s   |S r,   r"   ©r¿   rÀ   r   r"   r"   r#   r‡  Æ
  s    z_kron_return.<locals>.retc                 S   s   |  |j¡S r,   )rÍ  r=  r2  r"   r"   r#   r‡  Ë
  s    c                 S   s   |  | j¡S r,   ©rÍ  rº   r2  r"   r"   r#   r‡  Ñ
  s    c                 S   s   |  |j¡S r,   r3  r2  r"   r"   r#   r‡  Ö
  s    c                 S   s   |d S r¸   r"   r2  r"   r"   r#   r‡  Û
  s    )rd   r   re   ræ   r
   )r¿   rÀ   Za_is_arrZb_is_arrr‡  r"   r"   r#   Ú_kron_return¿
  s*    




r4  c                    sX   t | ddd� t |ddd� t| ƒ‰t|ƒ‰t| |ƒ‰t| d| ƒ‰ ‡ ‡‡‡fdd„}|S )NÚkronFr  r   c              	      sØ   ˆ| ƒ}ˆ|ƒ}|j d }|j d }|j d }|j d }|| }|| }	tj||	fˆ d�}
t|ƒD ]h}|| }t|ƒD ]R}|| }||d d …f }t|ƒD ],}|| }|||f | |
|||| …f< qšqvqbˆ| ||
ƒS )Nr8  r7  r:  )rº   r»   r½   rH  )r¿   rÀ   r#  r$  rŸ  r¥  r   r¦  ÚcmÚcnr_   rˆ  Zrjmpr±   ZirjmpÚslcr‰  Zcjmp©rÑ  Úfix_aÚfix_bZret_cr"   r#   Ú	kron_implï
  s$    



&zkron_impl.<locals>.kron_impl)r  r1  r4  r<   )r¿   rÀ   r<  r"   r9  r#   r<  á
  s    
(r<  )r   )F)T)T)r   )r»  )rË  )N)N)N)r   )T)N)ŸrL   Ú
contextlibrî   Úllvmliter   Únumpyr»   ÚoperatorÚnumba.core.imputilsr   r   r   r   Únumba.core.typingr   Únumba.core.extendingr   r	   r
   Ú
numba.corer   r   Únumba.core.errorsr   r   r   Úarrayobjr   r   r   Únumba.npr   rp  r…   r†   Ú
as_pointerZ	ll_char_prˆ   Úll_intcZ	ll_intc_pr‡   Z	ll_intp_pr¾  r;  rQ   Úfloat32r\   Ú	complex64Ú
complex128r   r$   r)   r+   r4   r5   rN   Úcontextmanagerro   rv   r‚   rƒ   r—   r¥   r·   rÅ   rÇ   rÉ   rÓ   r¾   rÙ   ÚmatmulrÚ   rÕ   rò   ró   rô   r  r
  r  rA   r=   r0  r  r$  r&  r'  r/  r3  r5  r  r6  r@  rA  rD  rE  rK  rL  rb  rc  rf  rg  rn  rs  rt  ru  r‚  r„  r‹  r‘  r•  r—  r™  r›  r¢  r£  r§  r¨  r´  r¶  rº  r¼  rÃ  rÆ  rÇ  rÈ  rÊ  rÌ  rÏ  r×  rØ  rÝ  rÞ  rà  râ  rè  ré  rð  rø  rê  rù  rú  rþ  r  r  r  r  r  r  r  r  r%  r&  r(  r)  r+  r1  r4  r5  r<  r"   r"   r"   r#   Ú<module>   s<  

    ü
) $
%2


,
9Y
#



	


'
	


+

r

v

0

0
A

P




$

^


<
 $$

:


F
# H

,

A

P



	"