
    Opf)                     z    d Z ddlZddlmZ ddlZg dZdZd Zd Z	d Z
d	 Zd
 Zd Zd Zd Zd Zd Zd Zd ZdS )zC
A functionally equivalent parser of the numpy.einsum input parser
    N)OrderedDict)
is_valid_einsum_charhas_valid_einsum_chars_only
get_symbolgen_unused_symbolsconvert_to_valid_einsum_charsalpha_canonicalizefind_output_strfind_output_shapepossibly_convert_to_numpyparse_einsum_input4abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZc                     | t           v p| dv S )u   Check if the character ``x`` is valid for numpy einsum.

    Examples
    --------
    >>> is_valid_einsum_char("a")
    True

    >>> is_valid_einsum_char("Ǵ")
    False
    z,->.)_einsum_symbols_basexs    Q/var/www/html/nettyfy-visnx/env/lib/python3.11/site-packages/opt_einsum/parser.pyr   r      s     %%71;7    c                 F    t          t          t          |                     S )u   Check if ``einsum_str`` contains only valid characters for numpy einsum.

    Examples
    --------
    >>> has_valid_einsum_chars_only("abAZ")
    True

    >>> has_valid_einsum_chars_only("Över")
    False
    )allmapr   )
einsum_strs    r   r   r   #   s     s'44555r   c                 L    | dk     rt           |          S t          | dz             S )u  Get the symbol corresponding to int ``i`` - runs through the usual 52
    letters before resorting to unicode characters, starting at ``chr(192)``.

    Examples
    --------
    >>> get_symbol(2)
    'c'

    >>> get_symbol(200)
    'Ŕ'

    >>> get_symbol(20000)
    '京'
    4      )r   chr)is    r   r   r   1   s(     	2vv#A&&q3w<<r   c              #   r   K   dx}}||k     r*t          |          }|dz  }|| v r|V  |dz  }||k     (dS dS )zGenerate ``n`` symbols that are not already in ``used``.

    Examples
    --------
    >>> list(oe.parser.gen_unused_symbols("abd", 2))
    ['c', 'e']
    r      Nr   )usednr   cntss        r   r   r   E   s_       KA
''qMM	Q99q ''''''r   c                     t          t          |           t          d          z
            }d t          |          D             d                    fd| D                       S )u  Convert the str ``einsum_str`` to contain only the alphabetic characters
    valid for numpy einsum. If there are too many symbols, let the backend
    throw an error.

    Examples
    --------
    >>> oe.parser.convert_to_valid_einsum_chars("Ĥěļļö")
    'cbdda'
    z,->c                 4    i | ]\  }}|t          |          S  r    ).0r   r   s      r   
<dictcomp>z1convert_to_valid_einsum_chars.<locals>.<dictcomp>b   s$    @@@TQ:a==@@@r    c              3   D   K   | ]}                     ||          V  d S Nget)r(   r   replacers     r   	<genexpr>z0convert_to_valid_einsum_chars.<locals>.<genexpr>c   s1      ::!8<<1%%::::::r   )sortedset	enumeratejoin)r   symbolsr/   s     @r   r   r   W   se     S__s5zz122G@@Yw-?-?@@@H77::::z::::::r   c                     t                      | D ]*}|dv r|vrt          t                              |<   +d                    fd| D                       S )u   Alpha convert an equation in an order-independent canonical way.

    Examples
    --------
    >>> oe.parser.alpha_canonicalize("dcba")
    'abcd'

    >>> oe.parser.alpha_canonicalize("Ĥěļļö")
    'abccd'
    z.,->r*   c              3   D   K   | ]}                     ||          V  d S r,   r-   )r(   r   renames     r   r0   z%alpha_canonicalize.<locals>.<genexpr>w   s1      666::a##666666r   )r   r   lenr4   )equationnamer8   s     @r   r	   r	   f   ss     ]]F 3 36>>v%c&kk22F4L776666X666666r   c                     |                      dd          d                    fdt          t                              D                       S )aU  
    Find the output string for the inputs ``subscripts`` under canonical einstein summation rules. That is, repeated indices are summed over by default.

    Examples
    --------
    >>> oe.parser.find_output_str("ab,bc")
    'ac'

    >>> oe.parser.find_output_str("a,b")
    'ab'

    >>> oe.parser.find_output_str("a,a,b,b")
    ''
    ,r*   c              3   N   K   | ]}                     |          d k    |V   dS )r   N)count)r(   r$   tmp_subscriptss     r   r0   z"find_output_str.<locals>.<genexpr>   s=      ZZ^=Q=QRS=T=TXY=Y=Y1=Y=Y=Y=YZZr   )replacer4   r1   r2   )
subscriptsr@   s    @r   r
   r
   z   sQ      ''R00N77ZZZZfS%8%899ZZZZZZr   c                 >     t           fd|D                       S )aO  Find the output shape for given inputs, shapes and output string, taking
    into account broadcasting.

    Examples
    --------
    >>> oe.parser.find_output_shape(["ab", "bc"], [(2, 3), (3, 4)], "ac")
    (2, 4)

    # Broadcasting is accounted for
    >>> oe.parser.find_output_shape(["a", "a"], [(4, ), (1, )], "a")
    (4,)
    c           	   3      K   | ]7t          d  t          fdD                       D                       V  8dS )c              3   8   K   | ]\  }}|d k    ||         V  dS )r   Nr'   )r(   shapelocs      r   r0   z.find_output_shape.<locals>.<genexpr>.<genexpr>   s6      ^^:5#UX\]U]U]E#JU]U]U]U]^^r   c                 :    g | ]}|                               S r'   )find)r(   r   cs     r   
<listcomp>z/find_output_shape.<locals>.<genexpr>.<listcomp>   s#    5P5P5PAaffQii5P5P5Pr   N)maxzip)r(   rJ   inputsshapess    @r   r0   z$find_output_shape.<locals>.<genexpr>   ss       p pcd^^V5P5P5P5P5P5P5P)Q)Q^^^^^p p p p p pr   )tuple)rN   rO   outputs   `` r   r   r      sN      p p p p phnp p p p p pr   c                 N    t          | d          st          j        |           S | S )aT  Convert things without a 'shape' to ndarrays, but leave everything else.

    Examples
    --------
    >>> oe.parser.possibly_convert_to_numpy(5)
    array(5)

    >>> oe.parser.possibly_convert_to_numpy([5, 3])
    array([5, 3])

    >>> oe.parser.possibly_convert_to_numpy(np.array([5, 3]))
    array([5, 3])

    # Any class with a shape is passed through
    >>> class Shape:
    ...     def __init__(self, shape):
    ...         self.shape = shape
    ...

    >>> myshape = Shape((5, 5))
    >>> oe.parser.possibly_convert_to_numpy(myshape)
    <__main__.Shape object at 0x10f850710>
    rF   )hasattrnp
asanyarrayr   s    r   r   r      s+    2 1g }Qr   c                 H    d}| D ]}|t           u r|dz  }|||         z  }|S )a  Convert user custom subscripts list to subscript string according to `symbol_map`.

    Examples
    --------
    >>>  oe.parser.convert_subscripts(['abc', 'def'], {'abc':'a', 'def':'b'})
    'ab'
    >>> oe.parser.convert_subscripts([Ellipsis, object], {object:'a'})
    '...a'
    r*   ...)Ellipsis)old_sub
symbol_mapnew_subr$   s       r   convert_subscriptsr\      sD     G % %==uGG z!}$GGNr   c                    t          |           }g }g }t          t          |           dz            D ]R}|                    |                    d                     |                    |                    d                     St          |          r|d         nd}d |D             } 	 t          t          j                            |                    }|	                    t                     d t          t          |                    D             n# t          $ r t          d          w xY wd                    fd	|D                       }||d
z  }|t          |          z  }|| fS )z:Convert 'interleaved' input to standard einsum input.
       r   Nc                 ,    g | ]}t          |          S r'   r   r(   r   s     r   rK   z-convert_interleaved_input.<locals>.<listcomp>   s!    CCC)!,,CCCr   c                 4    i | ]\  }}|t          |          S r'   r    )r(   idxsymbols      r   r)   z-convert_interleaved_input.<locals>.<dictcomp>   s$    ___+#vfjoo___r   ziFor this input type lists must contain either Ellipsis or hashable and comparable object (e.g. int, str).r=   c              3   8   K   | ]}t          |          V  d S r,   )r\   )r(   subrZ   s     r   r0   z,convert_interleaved_input.<locals>.<genexpr>   s.      XX#,S*==XXXXXXr   ->)listranger9   appendpopr2   	itertoolschainfrom_iterablediscardrX   r3   r1   	TypeErrorr4   r\   )	operandstmp_operandsoperand_listsubscript_listpoutput_list
symbol_setrB   rZ   s	           @r   convert_interleaved_inputry      s    >>LLN3x==A%&& 3 3L,,Q//000l..q112222&),&7&7A,r""TKCClCCCHN66~FFGG
 	8$$$ `_6R\K]K]A^A^___

 N N N M N N 	NN XXXXXXXXXJd
(jAAA
xs   -A,D D4c           	         t          |           dk    rt          d          t          | d         t                    r1| d                             dd          }d | dd         D             } nt          |           \  }} d|v sd	|v r\|                    d          dk    p|                    d	          dk    }|s|                    d
          dk    rt          d          d|v rn|                    dd                              dd                              d
d          }d                    t          |t          d | D                                           }d}d
|v r0|
                    d
          \  }}|
                    d          }d}	n|
                    d          }d}	t          |          D ]\  }
}d|v r|                    d          dk    s|                    d          dk    rt          d          | |
         j        dk    rd}n;t          t          | |
         j                  d          t          |          dz
  z
  }||k    r|}|dk     rt          d          |dk    r|                    dd          ||
<   |                    d|| d                   ||
<   d                    |          }|dk    rd}n|| d         }|	r|d
|                    d|          z   z  }nYt          |          }d                    t          t          |          t          |          z
                      }|d
|z   |z   z  }d
|v r|
                    d
          \  }}n|t          |          }}|D ](}||vr"t          d                    |                    )t          |
                    d                    t          |           k    rt          d          ||| fS )af  
    A reproduction of einsum c side einsum parsing in python.

    Returns
    -------
    input_strings : str
        Parsed input strings
    output_string : str
        Parsed output string
    operands : list of array_like
        The operands to use in the numpy contraction

    Examples
    --------
    The operand list is simplified to reduce printing:

    >>> a = np.random.rand(4, 4)
    >>> b = np.random.rand(4, 4, 4)
    >>> parse_einsum_input(('...a,...a->...', a, b))
    ('za,xza', 'xz', [a, b])

    >>> parse_einsum_input((a, [Ellipsis, 0], b, [Ellipsis, 0]))
    ('za,xza', 'xz', [a, b])
    r   zNo input operands r*   c                 ,    g | ]}t          |          S r'   ra   rb   s     r   rK   z&parse_einsum_input.<locals>.<listcomp>  s!    GGGQ-a00GGGr   r   N->rh   z%Subscripts can only contain one '->'..r=   c              3   >   K   | ]}t          |j                  V  d S r,   )r9   rF   rb   s     r   r0   z%parse_einsum_input.<locals>.<genexpr>#  s*      ;[;[QCLL;[;[;[;[;[;[r   TF   rW   zInvalid Ellipses.r'   zEllipses lengths do not match.z1Output character '{}' did not appear in the inputzDNumber of einsum subscripts must be equal to the number of operands.)r9   
ValueError
isinstancestrrA   ry   r?   r4   r   rL   splitr3   rF   r
   r1   r2   format)rr   rB   invalidr!   ellipse_indslongest	input_tmp
output_subsplit_subscriptsout_subnumrg   ellipse_countout_ellipseoutput_subscriptnormal_indsinput_subscriptschars                     r   r   r      sd   4 8}},---(1+s## Ca[((b11
GG(122,GGG  9BB
H 	zsj00##C((1,L*2B2B32G2G!2K 	Fz''--22DEEE j!!#r**223;;CCD"MMww1$;[;[RZ;[;[;[8[8[\\]] :$.$4$4T$:$:!Iz(s33GG)//44G!"233 	^ 	^HCczzIIcNNa''SYYu-=-=-B-B$%8999 C=&",,$%MM$'HSM,?(@(@!$D$DCST$UM 7**+G 1$$$%EFFF"a'',/KKr,B,B$S)),/KK|]NOO?\,],]$S)XX.//
 a<<KK&xyy1K 	;$!3!3E;!G!GGGJJ  /z::''&-=)>)>[AQAQ)Q"R"RSSK$,{::J z-7-=-=d-C-C***-79T9T* ! _ _'''PWWX\]]^^^ ( !!#&&''3x==88bccc-x77r   )__doc__rm   collectionsr   numpyrT   __all__r   r   r   r   r   r   r	   r
   r   r   r\   ry   r   r'   r   r   <module>r      s        # # # # # #       N 8 8 86 6 6  (  $; ; ;7 7 7([ [ [&p p p"  >  ("  "  " Jn8 n8 n8 n8 n8r   