o
    ¨i e?  ã                   @   sX   d Z ddlZddlZddlmZmZ dd„ Zdd„ Zdd	„ Z	d
d„ Z
dd„ Zdd„ ZdS )zf
Tets a series of opt_einsum contraction paths to ensure the results are the same for different paths
é    N)ÚcontractÚcontract_expressionc               	   C   s  t  t¡ tddddd W d   ƒ n1 sw   Y  t  t¡ tddddƒ W d   ƒ n1 s4w   Y  t d¡} t  t¡ tddd| d W d   ƒ n1 sVw   Y  tdddƒ}t  t¡}|tj d	d
¡ƒ W d   ƒ n1 szw   Y  dt|j	ƒv sˆJ ‚t  t¡}|tj d	d
¡tj d	d
¡tj d	d
¡ƒ W d   ƒ n1 s­w   Y  dt|j	ƒv s»J ‚t  t¡}|tj d	d
d¡tj d
d¡ƒ W d   ƒ n1 sÛw   Y  dt|j	ƒv séJ ‚t  t¡}|tj d	d¡tj d
dd¡ƒ W d   ƒ n	1 s
w   Y  dt|j	ƒv sJ ‚t  t¡}|tj d	d
¡tj d
d¡tj d	dd¡d W d   ƒ n	1 sAw   Y  dt|j	ƒv sPJ ‚t  t¡}|tj d	d
¡tj d
d¡dd W d   ƒ n	1 srw   Y  dt|j	ƒv sJ ‚d S )Nz	ab,bc->ac)é   é   )r   é   F©Úoptimize)é*   r	   )r   r   )Úoutr   r   z$`ContractExpression` takes exactly 2r   z4Internal error while evaluating `ContractExpression`é   é   ÚF)Úorderz6only valid keyword arguments to a `ContractExpression`)
ÚpytestÚraisesÚ
ValueErrorr   ÚnpÚemptyÚrandomÚrandÚstrÚvalue)r
   ÚexprÚerr© r   úO/var/www/myenv/lib/python3.10/site-packages/opt_einsum/tests/test_edge_cases.pyÚtest_contract_expression_checks   sF   ÿÿ
ÿÿ,ÿ"ÿ"ÿ0ÿ$ÿr   c            	      C   sÀ   t j ddd¡} t j dd¡}t j dd¡}t j d¡}td| ||dd}td| ||d	d}t  ||¡s6J ‚|| }td
| |||dd}td
| |||d	d}t  ||¡sVJ ‚t  ||¡s^J ‚d S )Né   r   r   r   é
   z	ijk,kl,jlFr   Tzijk,kl,jl,i->i©r   r   r   r   Úallclose©	ÚaÚbÚcÚdÚ
ein_scalarÚ
opt_scalarÚresultÚeinÚoptr   r   r   Útest_broadcasting_contraction7   s   r+   c            	      C   sÄ   t j dddd¡} t j dd¡}t j dd¡}t j dd¡}td| ||dd}td| ||d	d}t  ||¡s8J ‚|| }td
| |||dd}td
| |||d	d}t  ||¡sXJ ‚t  ||¡s`J ‚d S )Nr   r   r   r   é   z
abjk,kl,jlFr   Tzabjk,kl,jl,ab->abr   r!   r   r   r   Útest_broadcasting_contraction2K   s   r-   c                  C   sx   t j ddd¡} t j ddd¡}t j dd¡}t j dd¡}td| |||dd}td| |||d	d}t  ||¡s:J ‚d S )
Nr   r   r   r   r,   zajk,kbl,jl,ab->abFr   Tr   )r"   r#   r$   r%   r)   r*   r   r   r   Útest_broadcasting_contraction3_   s   r.   c                  C   sH   t  d¡ ddd¡} td| | dd}td| | dd}t  ||¡s"J ‚d S )	Né@   r   r   é   zobk,ijk->iojFr   T)r   ÚarangeÚreshaper   r    )r"   r)   r*   r   r   r   Útest_broadcasting_contraction4l   s   r3   c                  C   sj   t ddddƒ} | jd d dksJ ‚| jd d d	u sJ ‚| jd
 d dks(J ‚| jd
 d dks3J ‚d S )Nzab,bc,bd->acd)r   r   )r   r   )r   é   r   r   z
bc,ab->bcaéÿÿÿÿFr   zbca,bd->acdÚGEMM)r   Úcontraction_list)r   r   r   r   Ú,test_can_blas_on_healed_broadcast_dimensionsu   s
   r8   )Ú__doc__Únumpyr   r   Ú
opt_einsumr   r   r   r+   r-   r.   r3   r8   r   r   r   r   Ú<module>   s    ,	