a
    
d0                     @   s  d dl mZ d dlmZ d dlZd dlZd dlmZm	Z	 d dl
mZmZmZmZ d dlmZ d dlmZ d dlmZmZ d d	lmZmZmZmZ e	 Zd
e_e	 Zd
e_dd Zdd Z dd Z!dd Z"dd Z#dd Z$dd Z%dd Z&dd Z'dd Z(dd  Z)d!d" Z*d#d$ Z+d%d& Z,d'd( Z-d)d* Z.d+d, Z/d-d. Z0d/d0 Z1d1d2 Z2d3d4 Z3d5d6 Z4d7d8 Z5d9d: Z6d;d< Z7edOd=d>Z8d?d@ Z9dAdB Z:dPdDdEZ;dFdG Z<dHdI Z=dJdK Z>G dLdM dMeeZ?e@dNkreA  dS )Q    )partial)permutationsN)compile_isolatedFlags)jitnjit
from_dtypetypeof)numpy_version)TypingError)typeserrors)TestCaseMemoryLeakMixinCompilationCachetagTc                    s    fdd}|S )aP  Decorator for generic check functions.
        Iterates over 'pyfuncs_to_use', calling 'func' with the iterated
        item as first argument. Example:

        @from_generic(numpy_array_reshape, array_reshape)
        def check_only_shape(pyfunc, arr, shape, expected_shape):
            # Only check Numba result to avoid Numpy bugs
            self.memory_leak_setup()
            got = generic_run(pyfunc, arr, shape)
            self.assertEqual(got.shape, expected_shape)
            self.assertEqual(got.size, arr.size)
            del got
            self.memory_leak_teardown()
    c                    s    fdd}|S )Nc                     s    fddD S )Nc                    s"   g | ]}|g R i qS  r   ).0pyfunc)argsfunckwargsr   l/var/www/html/stable-diffusion-webui/venv/lib/python3.9/site-packages/numba/tests/test_array_manipulation.py
<listcomp>'       zCfrom_generic.<locals>.decorator.<locals>.result.<locals>.<listcomp>r   r   r   )r   pyfuncs_to_user   r   result&   s    z/from_generic.<locals>.decorator.<locals>.resultr   )r   r   r   )r   r   	decorator%   s    zfrom_generic.<locals>.decoratorr   )r   r   r   r   r   from_generic   s    r    c                 C   s
   |  |S N)reshapearrZnewshaper   r   r   array_reshape,   s    r%   c                 C   s   t | |S r!   )npr"   r#   r   r   r   numpy_array_reshape0   s    r'   c                 C   s   t | |S r!   r&   broadcast_to)r$   shaper   r   r   numpy_broadcast_to4   s    r+   c                  G   s
   t j|  S r!   )r&   Zbroadcast_shapesr   r   r   r   numpy_broadcast_shapes8   s    r-   c                  G   s
   t j|  S r!   )r&   Zbroadcast_arraysr,   r   r   r   numpy_broadcast_arrays<   s    r.   c                 C   s   t | || S r!   r(   )r$   r*   idxr   r   r   numpy_broadcast_to_indexing@   s    r0   c                 C   s   |   S r!   )flattenar   r   r   flatten_arrayD   s    r4   c                 C   s   |   S r!   ravelr2   r   r   r   ravel_arrayH   s    r7   c                 C   s
   |   jS r!   )r6   sizer2   r   r   r   ravel_array_sizeL   s    r9   c                 C   s
   t | S r!   )r&   r6   r2   r   r   r   numpy_ravel_arrayP   s    r:   c                 C   s   |   S r!   	transposer2   r   r   r   transpose_arrayT   s    r=   c                 C   s
   t | S r!   r&   r<   r2   r   r   r   numpy_transpose_arrayX   s    r?   c                 C   s   t j| |dS N)axesr>   r$   rA   r   r   r    numpy_transpose_array_axes_kwarg[   s    rC   c                 C   s   t j| |d S r@   )r&   r<   copyrB   r   r   r   %numpy_transpose_array_axes_kwarg_copy_   s    rE   c                 C   s
   |  |S r!   r;   rB   r   r   r   array_transpose_axesc   s    rF   c                 C   s   |  | S r!   )r<   rD   rB   r   r   r   array_transpose_axes_copyg   s    rG   c                 C   s^   t t | | d | d|f}t t |d |df}||j j}t |d}|d }|S )N   )   r      rJ   )r&   r"   arangeTr<   )mnZr1Zr2Zr_difzr   r   r   transpose_issue_4708k   s     rP   c                 C   s   |   S r!   )squeezer2   r   r   r   squeeze_arrayt   s    rR   c                 C   s   t | |S r!   )r&   expand_dims)r3   axisr   r   r   rS   x   s    rS   c                  G   s
   t j|  S r!   )r&   
atleast_1dr,   r   r   r   rU   |   s    rU   c                  G   s
   t j|  S r!   )r&   
atleast_2dr,   r   r   r   rV      s    rV   c                  G   s
   t j|  S r!   )r&   
atleast_3dr,   r   r   r   rW      s    rW   c                 C   s0   | j d d f| j dd   }tjjj| |dS )Nr   rI   rJ   )strides)rX   r&   libstride_tricks
as_strided)r3   rX   r   r   r   as_strided1   s    r\   c                 C   sN   d}| j d d | j d | d |f }| j| jd f }tjjj| ||dS )NrH   rJ   )r*   rX   )r*   rX   r&   rY   rZ   r[   )r3   Zwindowr*   rX   r   r   r   as_strided2   s    $r^   c                 C   s   t jjj| ||dS )N)rT   )r&   rY   rZ   sliding_window_view)xwindow_shaperT   r   r   r   r_      s    r_   c                 C   s   | j f}| j}d|||f< d S )N      ?)r`   y)r$   Zarr2dr`   rc   r   r   r   	bad_index   s    rd   c                 C   s   | d S )N)rJ   g       @r   r$   r   r   r   bad_float_index   s    rf   Fc                 C   s   t | ||S r!   )r&   Zfill_diagonal)r$   valwrapr   r   r   numpy_fill_diagonal   s    ri   c                 C   s
   t | S r!   )r&   r*   re   r   r   r   numpy_shape   s    rj   c                 C   s
   t | S r!   )r&   Zflatnonzeror2   r   r   r   numpy_flatnonzero   s    rk   c                 C   s
   t | S r!   )r&   Zargwherer2   r   r   r   numpy_argwhere   s    rl   c                       s  e Zd ZdZ fddZdd Zdd Zdd	 Zd
d Zdd Z	dd Z
dd Zdd Zdd Zdd Zdd Zdd ZedfddZdd  Zefd!d"Zefd#d$Zd%d& Zd'd( Zefd)d*Zd+d, Zefd-d.Zd/d0 Zd1d2 Zd3d4 Zd5d6 Zd7d8 Zd9d: Z d;d< Z!d=d> Z"d?d@ Z#dAdB Z$dCdD Z%dEdF Z&dGdH Z'dIdJ Z(dKdL Z)dMdN Z*dOdP Z+dQdR Z,dSdT Z-dUdV Z.e/dWdX Z0dYdZ Z1d[d\ Z2d]d^ Z3d_d` Z4dadb Z5dcdd Z6dedf Z7dgdh Z8didj Z9dkdl Z:dmdn Z;dodp Z<dqdr Z=dsdt Z>dudv Z?  Z@S )wTestArrayManipulationz4
    Check shape-changing operations on arrays.
    c                    s   t t|   t | _d S r!   )superrm   setUpr   ccacheself	__class__r   r   ro      s    zTestArrayManipulation.setUpc                    s4  t tg}dd t|fdd t|fddt|fdd}t|fd	d
t|fdd}td} |d  |d  |d  |d  |d tdd} |d  |d  |d  |d  |d |d |d |d tdd} |d  |d  |d  |d tddj} |d  |d ||d ||d ||d tddd} |d  |d  |d  |d   |d! |d" ||d# ||d$  fd%d&}tg }|| ||d'   d S )(Nc                 S   s"   t | t|t|f}|||S r!   )r   r	   entry_point)r   r$   r*   cresr   r   r   generic_run   s    z=TestArrayManipulation.test_array_reshape.<locals>.generic_runc                    s8   | ||}    | ||}|| ~  d S r!   memory_leak_setupassertPreciseEqualmemory_leak_teardown)r   r$   r*   expectedgotrw   rr   r   r   check   s    
z7TestArrayManipulation.test_array_reshape.<locals>.checkc                    s@       | ||}|j| |j|j ~  d S r!   )ry   assertEqualr*   r8   r{   )r   r$   r*   expected_shaper}   r~   r   r   check_only_shape   s    zBTestArrayManipulation.test_array_reshape.<locals>.check_only_shapec                    sL    t} | || W d    n1 s,0    Y  t|jd d S )Nzincompatible shape for array)assertRaisesNotImplementedErrorr   str	exceptionr   r$   r*   raisesr~   r   r   check_err_shape   s
    *zATestArrayManipulation.test_array_reshape.<locals>.check_err_shapec                    sL    t} | || W d    n1 s,0    Y  t|jd d S )Nz)total size of new array must be unchangedr   
ValueErrorr   r   r   r   r~   r   r   check_err_size   s
    *z@TestArrayManipulation.test_array_reshape.<locals>.check_err_sizec                    sL    t} | || W d    n1 s,0    Y  t|jd d S )Nzmultiple negative shape valuesr   r   r~   r   r   check_err_multiple_negative   s
    *zMTestArrayManipulation.test_array_reshape.<locals>.check_err_multiple_negative   )r   )      )   rH   )r   rJ   rH   )rJ   r   rJ   rJ   rH   rJ   rI   rH   r   r   )   )r   r   )r   rH   rI   )rJ   r   rJ   rH   rJ   rI   rJ   r   r   )rI      r      r]   r]   )r]   r   )r   r]   r   )r   r   r]   )r]   r   )r]   r   r   )r   r   r]   r]   c                    s    | d  | d  | d  | d | dd | dd | dd | d	d
 | dd | d | d | d d S )Nr   r   )rJ   r   rI   )r   7   rJ   r   rI   r]   r   r   r]   r   r   )r   r]   )r   r   )r]   r   r   )r   r   r   r   rJ   rJ   rI   r   re   )r   r   r   r   r   check_empty  s    





z=TestArrayManipulation.test_array_reshape.<locals>.check_empty)rH   rI   r   )	r%   r'   r    r&   rK   r"   rL   arraydisable_leak_check)rr   r   r   r   r$   r   r   )r   r   r   rw   rr   r   test_array_reshape   sh    	































z(TestArrayManipulation.test_array_reshapec              	      s  t tttg}fddt|fdd}t|fdd}t|fdd}t|fd	d
}tdtdddtddddtdddddtddddd d dd d dd d f g}tt	|D ]j}||| d  t
tt|| jD ]@}t	| t fdd|D }	||| | ||| |	 qqttgdd }
|
dd |
dd |
dd   ||d d ||d d ||d d ||d d ||d d ||d d ||d d ||d d  ||d d! ||d d" ||d d# ||d d$ ||d d% ||d d& ||d d' ||d d( t.}td)d*ttddgf W d    n1 s0    Y  d+t|j d S ),Nc                    s&    j | t|t|f}|||S r!   rp   compiler	   ru   )r   r$   rA   rv   rq   r   r   run7  s    z<TestArrayManipulation.test_array_transpose_axes.<locals>.runc                    sN   | ||} | ||} || |jj|jj |jj|jj d S r!   )rz   r   flagsf_contiguousc_contiguous)r   r$   rA   r|   r}   r   rr   r   r   r   ;  s    


z>TestArrayManipulation.test_array_transpose_axes.<locals>.checkc                    sL    t} | || W d    n1 s,0    Y  t|jd d S )Nzrepeated axis in transposer   r   r$   rA   r   r   r   r   check_err_axis_repeatedE  s
    *zPTestArrayManipulation.test_array_transpose_axes.<locals>.check_err_axis_repeatedc                    sL    t} | || W d    n1 s,0    Y  t|jd d S )Nz2axis is out of bounds for array of given dimensionr   r   r   r   r   check_err_axis_oobL  s
    *zKTestArrayManipulation.test_array_transpose_axes.<locals>.check_err_axis_oobc                    s>    ttf  | || W d    n1 s00    Y  d S r!   )r   	TypeErrorr   )r   r$   rA   r   r   r   check_err_invalid_argsS  s    zOTestArrayManipulation.test_array_transpose_axes.<locals>.check_err_invalid_argsr   r   r   rI   rH   rJ   @   r   c                    s   g | ]}|  qS r   r   r   r`   )ndimr   r   r   d  r   zCTestArrayManipulation.test_array_transpose_axes.<locals>.<listcomp>c                 S   s*   | ||}t | ||}tj|| d S r!   )r   r&   testingassert_equal)r   rM   rN   r|   r}   r   r   r   check_issue_4708h  s    
zITestArrayManipulation.test_array_transpose_axes.<locals>.check_issue_4708r   foo)r   g333333@)333333?r   r   )rI   r   r   )rH   rI   rJ   rJ   r   rJ   )r   r   rI   )rI   r   r   )rH   rJ   rI   )rJ   rI   )rH   rJ   rI   r   )rH   rJ   rI   TZnopythonz#np.transpose does not accept tuples)rC   rE   rF   rG   r    r&   rK   r"   rangelenr   tupler   rP   r   r   r   r   r?   r   assertInr   r   )rr   r   r   r   r   r   ZarrsirA   Zneg_axesr   er   )r   r   rr   r   test_array_transpose_axes1  sl    	,



>z/TestArrayManipulation.test_array_transpose_axesc                    s   t fddfdd  fdd}td}|| tdd	}|| ||j ||d d d
  td}|| d S )Nc                    s&   j  t| t|f}|| |S r!   r   )r$   rT   rv   r   rr   r   r   r     s    z3TestArrayManipulation.test_expand_dims.<locals>.runc                    s6    | |}   | |}|| ~  d S r!   rx   )r$   rT   r|   r}   r   r   rr   r   r   r     s    

z5TestArrayManipulation.test_expand_dims.<locals>.checkc                    s,   t | j d | jd D ]} | | qd S )NrJ   )r   r   )r$   rT   )r   r   r   check_all_axes  s    z>TestArrayManipulation.test_expand_dims.<locals>.check_all_axesr   r   r   r]   *   )rS   r&   rK   r"   rL   r   )rr   r   r$   r   )r   r   r   rr   r   test_expand_dims  s    


z&TestArrayManipulation.test_expand_dimsc                 C   s   t }tdd|}td}|  }|dd W d    n1 sD0    Y  | dt|j |  }||d W d    n1 s0    Y  | dt|j d S )NTr   r   hellorH   z#First argument "a" must be an arrayz"Argument "axis" must be an integer)rS   r   r&   rK   assertTypingErrorr   r   r   )rr   r   cfuncr$   r   r   r   r   test_expand_dims_exceptions  s    

(
(z1TestArrayManipulation.test_expand_dims_exceptionsc                    s(  fdd fdd} fdd}t d}t d}|| ||| t d	}t d
d }|d d d }	|| ||||	 t d
d}
|
j}|
d d d }||
 ||
|| t dd}|j}|d d d }|| |||| t dd}|| ||||| d S )Nc                    s$     | |  |  |  d S r!   )ZassertStridesEqualrz   r1   )r}   r|   rq   r   r   check_result  s    z<TestArrayManipulation.check_atleast_nd.<locals>.check_resultc                    s    | |  d S r!   r   )arg)r   r   r   r   r   check_single  s    z<TestArrayManipulation.check_atleast_nd.<locals>.check_singlec                     sF   |  } |  } t|t| t||D ]\}}|| q.d S r!   )r   r   zip)r   Zexpected_tupleZ	got_tupler}   r|   r   r   r   rr   r   r   check_tuple  s
    z;TestArrayManipulation.check_atleast_nd.<locals>.check_tupler   y              @r   r                 ?r]   rI   rH   r   r      )rI   rI   rI   rI   )r&   r   rK   r"   rL   )rr   r   r   r   r   Za1Za2b1b2Zb3Zc1c2c3d1Zd2Zd3r   r   r   r   check_atleast_nd  s2    



z&TestArrayManipulation.check_atleast_ndc                 C   s"   t }tdd|}| || d S NTr   )rU   r   r   rr   r   r   r   r   r   test_atleast_1d  s    z%TestArrayManipulation.test_atleast_1dc                 C   s"   t }tdd|}| || d S r   )rV   r   r   r   r   r   r   test_atleast_2d  s    z%TestArrayManipulation.test_atleast_2dc                 C   s"   t }tdd|}| || d S r   )rW   r   r   r   r   r   r   test_atleast_3d  s    z%TestArrayManipulation.test_atleast_3dc                    sP    fdd fdd}t d}|| ||d ||d d S )Nc                    s   j  t| f}|| S r!   r   )r$   rv   r   r   r   r     s    z3TestArrayManipulation.check_as_strided.<locals>.runc                    s     | }| } || d S r!   rz   )r$   r|   r}   r   r   r   r     s    z5TestArrayManipulation.check_as_strided.<locals>.checkr   r   )r   rJ   r   )r&   rK   r"   )rr   r   r   r$   r   r   r   check_as_strided  s    
z&TestArrayManipulation.check_as_stridedc                 C   s   |  t |  t d S r!   )r   r\   r^   rq   r   r   r   test_as_strided  s    
z%TestArrayManipulation.test_as_stridedc              	      s   fdd}t d}dD ]@} d|  ||d| W d    q1 sP0    Y  qt ddd	}d
D ]@} d|  ||d| W d    qr1 s0    Y  qrdD ]@} d|  ||d| W d    q1 s0    Y  qt ddddd}dD ]D} d|  ||d| W d    n1 sL0    Y  q d ||dd W d    n1 s0    Y  d S )Nc                    s2   t jjj| ||dd}t| ||} || d S )NT)	writeable)r&   rY   rZ   r_   rz   )r$   ra   rT   r|   r}   rq   r   r   r     s
    z=TestArrayManipulation.test_sliding_window_view.<locals>.checkr   )Nr   r]   r   z1d array, axis=r      
      r   z2d array, axis=)Nr   rJ   rJ   r   )rJ   r   )r   r   r   rI   )r   )r   r   z4d array, axis=rH   rI   z2d array, repeated axes)r   rH   rH   )r   rJ   r   )r&   rK   subTestr"   )rr   r   arr1rT   arr2Zarr4r   rq   r   test_sliding_window_view  s$    	
,,,0z.TestArrayManipulation.test_sliding_window_viewc                    s   fdd} fdd}    td}tddd} d	 |d
|dd  W d    n1 sj0    Y   d |d
|dd  W d    n1 s0    Y   d |d|dd  W d    n1 s0    Y   d |d|dd  W d    n1 s0    Y   d |d|dd  W d    n1 sN0    Y   d |d|dd W d    n1 s0    Y   d |d|dd W d    n1 s0    Y   d |d|dd W d    n1 s0    Y   d |d|dd W d    n1 s60    Y   d  |d|d!d" W d    n1 sp0    Y   d# |d|d!d$ W d    n1 s0    Y   d% |d&|d  W d    n1 s0    Y   d' |d&|d( W d    n1 s0    Y   d) |d&|d* W d    n1 sR0    Y   d+ |d,|dd( W d    n1 s0    Y   d- |d,|dd* W d    n1 s0    Y  d S ).Nc                    sH     t}t|  W d    n1 s(0    Y   | t|j d S r!   )r   r   r_   r   r   r   msgr   r   rq   r   r   _raises1  s    &zFTestArrayManipulation.test_sliding_window_view_errors.<locals>._raisesc                    sJ     tj}t|  W d    n1 s*0    Y   | t|j d S r!   )r   r   r   r_   r   r   r   r   rq   r   r   _typing_error6  s    &zLTestArrayManipulation.test_sliding_window_view_errors.<locals>._typing_errorr   r   r   r   z1d window shape too largezwindow_shape cannot be largerr   z2d window shape too large)r      z1d window negative sizez&`window_shape` cannot contain negativer]   z2d window with a negative size)r   r   z1d array, 2d window shapez%matching length window_shape and axis)r   rI   z$2d window shape, only one axis givenrJ   z1d window shape, 2 axes givenr   r   z1d array, second axiszArgument axis out of boundsr   z1d array, axis -2r   z2d array, fourth axis)r   r   r   rH   z2d array, axis -3r   zwindow_shape=Nonez3window_shape must be an integer or tuple of integerzwindow_shape=float@zwindow_shape=tuple(float))r   z
axis=floatz1axis must be None, an integer or tuple of integerzaxis=tuple(float))r   r&   rK   r"   r   )rr   r   r   r   r   r   rq   r   test_sliding_window_view_errors0  s^    
,,,........$$$$z5TestArrayManipulation.test_sliding_window_view_errorsCc           
      C   s   t ddd}|dkr |j}t}t|}|dkr@|jdd}| |j| t	||f|d}|j
}||}||}	t j||	 d S )N	   rH   FA)layoutr   )r&   rK   r"   rL   r4   r	   rD   r   r   r   ru   r   r   )
rr   r   r   r3   r   
arraytype1crr   r|   r}   r   r   r   test_flatten_arrayt  s    z(TestArrayManipulation.test_flatten_arrayc                 C   s,   | j td | j tdd | j tdd d S )Nr   r   )r   r   r   )r   no_pyobj_flagsrq   r   r   r   test_flatten_array_npm  s    z,TestArrayManipulation.test_flatten_array_npmc                    s  fdd}t |tt |t  fdd}|tddddd |tdjddd	d
d	d |tddddd d d d df dd |tdddddd |tdjdddd	d
d	d |tdddddd d d d d d df dd d S )Nc           
         s   t |}|j| t| |f d}|j}| |}||}tj|| |jj	|jj	k}|jj	|jj	k}	||dk ||	 d S )Nr   r   )
r	   r   r   r   ru   r&   r   r   ctypesdata)
r   r3   assume_layoutr   r   r   r|   r}   Z	py_copiedZ	nb_copied)r   rr   r   r   generic_check  s    z=TestArrayManipulation.test_ravel_array.<locals>.generic_checkc                     s    | i |  | i | d S r!   r   r   )check_functioncheck_methodr   r   r     s    z5TestArrayManipulation.test_ravel_array.<locals>.checkr   rH   r   )r  r   order   rI   r   r   $   )r   r7   r:   r&   rK   r"   )rr   r   r  r   r   )r  r  r   rr   r   test_ravel_array  s    

. z&TestArrayManipulation.test_ravel_arrayc           	      C   sV   t ddd}t}t|}t||f|d}|j}||}||}t j|| d S Nr   rH   r   )	r&   rK   r"   r9   r	   r   ru   r   r   	rr   r   r3   r   r   r   r   r|   r}   r   r   r   test_ravel_array_size  s    z+TestArrayManipulation.test_ravel_array_sizec                 C   s   | j td d S Nr   )r  r   rq   r   r   r   test_ravel_array_npm  s    z*TestArrayManipulation.test_ravel_array_npmc                 C   s   | j td d S r  )r  r   rq   r   r   r   test_ravel_array_size_npm  s    z/TestArrayManipulation.test_ravel_array_size_npmc                    s"   t ttg fdd}|  d S )Nc                    sR   t ddd}t|}t| |f d}|j}| |}||}t j|| d S r  )r&   rK   r"   r	   r   ru   r   r   )r   r3   r   r   r   r|   r}   r   r   r   r     s    z9TestArrayManipulation.test_transpose_array.<locals>.check)r    r=   r?   )rr   r   r   r   r   r   test_transpose_array  s    
z*TestArrayManipulation.test_transpose_arrayc                 C   s   | j td d S r  )r  r   rq   r   r   r   test_transpose_array_npm  s    z.TestArrayManipulation.test_transpose_array_npmc           	      C   s\   t dddddd}t}t|}t||f|d}|j}||}||}t j|| d S )Nr   rI   rJ   rH   r   r   )	r&   rK   r"   rR   r	   r   ru   r   r   r  r   r   r   test_squeeze_array  s    z(TestArrayManipulation.test_squeeze_arrayc                 C   sN   |  tj}| jtd W d    n1 s.0    Y  | dt|j d S )Nr   rQ   )r   r   r   r  r   r   r   r   rr   r   r   r   r   test_squeeze_array_npm  s    *z,TestArrayManipulation.test_squeeze_array_npmc                 C   sh  t dd }t dd }tdd}td tjftjtd ftdtjdftjdtd ftdttjdfdtjtftjtdtjdfdtd tjftjtdttjdfdtjtjtftjdtjtftd	dtjd ftjdtd fg}|j}|}|j}|}|D ]f}	|||	}
|||	}tj	|
| t
|}||	 }|| |	|}
|| |	|}tj	|
| qd S )
Nc                 S   s   | | S r!   r   )r3   r/   r   r   r   np_new_axis_getitem  s    z@TestArrayManipulation.test_add_axis.<locals>.np_new_axis_getitemc                 S   s   || |< | S r!   r   )r3   r/   itemr   r   r   np_new_axis_setitem  s    z@TestArrayManipulation.test_add_axis.<locals>.np_new_axis_setitemiH  )r   r   r      rJ   rI   rH   )r   r&   rK   r"   sliceZnewaxisEllipsispy_funcr   r   Z
zeros_likerD   )rr   r  r  r3   Z	idx_casesZpyfunc_getitemZcfunc_getitemZpyfunc_setitemZcfunc_setitemr/   r|   r}   Za_emptyr  r   r   r   test_add_axis  s@    





z#TestArrayManipulation.test_add_axisc                 C   s~   |   P}ttdtjfdtjfg}ttjdd}tt||ft	d W d    n1 s^0    Y  | 
dt|j d S )Nr`   rc   rI   r   r   zUnsupported array index type)r   r   r&   dtypeint32r   Arrayr   rd   r   r   r   r   )rr   r   r   Z
arraytype2r   r   r   test_bad_index_npm  s    

$z(TestArrayManipulation.test_bad_index_npmc                 C   sV   |   (}ttttjddf W d    n1 s60    Y  | dt|j d S )NrI   r   z$Unsupported array index type float64)	r   r   rf   r   r"  float64r   r   r   r  r   r   r   test_bad_float_index_npm  s    
"z.TestArrayManipulation.test_bad_float_index_npmc                    s   t tdddd  dd } fdd} fd	d
}fdd}|dD ]}| D ]}||| q\qR|dD ]}| D ]}||| qqvtdgd dd}||d ||g d ||tg d d S )NTr   c                 s   sX   | | fV  d|  | fV  | d|  fV  d|  d d|  d fV  | | | | fV  dV  d S )NrI   rJ   rJ   rJ   rJ   r   )rN   r   r   r   _shape_variations%  s    
zITestArrayManipulation.test_fill_diagonal_basic.<locals>._shape_variationsc                   s   s   dV  dV  t jV  t j V  dV  t dV  dV  ddgV  t ddd	d
dV  t t dd	d	V  t dd	d	d d d V  d S )NrJ   gA`"	@Tr   r   r   r   6   rH   rI   r]   )r&   naninfrK   r"   asfortranarrayr   r   r   r   _val_variations3  s    

zGTestArrayManipulation.test_fill_diagonal_basic.<locals>._val_variationsc                 3   s<    | D ].}t j|t jdV  t t j|t jdV  qd S )Nr   )r&   zerosr$  r,  ones)rN   r*   r'  r   r   #_multi_dimensional_array_variationsA  s    z[TestArrayManipulation.test_fill_diagonal_basic.<locals>._multi_dimensional_array_variationsc                 3   sL    | D ]>}t jtdd |D t jd}tdd |D }|| V  qd S )Nc                 S   s   g | ]}|d  qS rI   r   r   r   r   r   r   I  r   zwTestArrayManipulation.test_fill_diagonal_basic.<locals>._multi_dimensional_array_variations_strided.<locals>.<listcomp>r.  c                 s   s   | ]}t d |d dV  qdS )r   rI   N)r  r   r   r   r   	<genexpr>J  r   zvTestArrayManipulation.test_fill_diagonal_basic.<locals>._multi_dimensional_array_variations_strided.<locals>.<genexpr>)r&   r/  r   r$  )rN   r*   tmpZslicerr1  r   r   +_multi_dimensional_array_variations_stridedG  s    zcTestArrayManipulation.test_fill_diagonal_basic.<locals>._multi_dimensional_array_variations_stridedc                    sd   dD ]Z}|   }|   }|d u r&i }nd|i}||fi |  ||fi | || qd S )N)NTFrh   )rD   rz   )r$   rg   rh   r3   bparamsr   r   rr   r   r   _check_fill_diagonalM  s    zLTestArrayManipulation.test_fill_diagonal_basic.<locals>._check_fill_diagonalrH   r   F)FTFTFT)ri   r   r&   r   r"   )rr   r-  r2  r6  r:  r$   rg   r   )r'  r   r   rr   r   test_fill_diagonal_basic!  s"    


z.TestArrayManipulation.test_fill_diagonal_basicc              	      sx  t }tdd| d}  tg tdfD ]J}t} || W d    n1 s`0    Y  dt|j	v s2J q2t
6}td} || dt|j	 W d    n1 s0    Y   fdd	}tjd
tjd}tj}||| dtjg}||| tjg dtjd}||| tjd
tjd}g d}||| d}||| tdg}||| d S )NTr   rJ   r   z'The first argument must be at least 2-D)rH   rH   r   z/All dimensions of input must be of equal lengthc                    sJ    t} | | W d    n1 s*0    Y  dt|j d S )Nz'Unable to safely conform val to a.dtyper   )r$   rg   r   r   rr   r   r   _assert_raises~  s    (zPTestArrayManipulation.test_fill_diagonal_exception_cases.<locals>._assert_raisesrH   rH   r.  ffffff
@)rJ   rI   g    _B)gffffff?g@}Ô%Igp=TrA  )ri   r   r   r&   r   r0  r   r   r   r   r   r/  r   r!  r*  r+  int64float32)rr   r   rg   r3   r   r>  r$   r   r=  r   "test_fill_diagonal_exception_casesi  s6    (

0





z8TestArrayManipulation.test_fill_diagonal_exception_casesc                 C   s  t }tdd|}tddgtddgtddgtddgtddgtdd	gtd
dgtd
dgtd
dgtddgtddgtddgtddgtddgtddgddgddgg}|D ](\}}|||}|||}| || qd S )NTr   r   r   r   rH   rJ   r3  rJ   rI   rH   rH   rJ   rH   r   rI   r   r   rI   rJ   )rI   r   rI   rI   )r+   r   r&   r   r0  rK   rz   )rr   r   r   r  input_arrayr*   r|   r}   r   r   r   test_broadcast_to  s0    

z'TestArrayManipulation.test_broadcast_toc                 C   sd   t }tdd|}tdddg}d}|D ]4}|||}|||}| || | |jd  q*d S )NTr   {   r   Z	WRITEABLE)r+   r   r&   r   rz   ZassertFalser   )rr   r   r   inputsr*   r$   r|   r}   r   r   r   test_broadcast_to_0d_array  s    


z0TestArrayManipulation.test_broadcast_to_0d_arrayc           	      C   sv  t }tdd|}tddtdgtddtdgtddtdgddtdgddtdgtddtd	gtdd
td	gtddtd	gtddtd	gtddtdgtddtdgtddtdgtddtdgtddtdgtddtdgddtdgddtdgg}|   |D ]V\}}}}| |}||| W d    n1 sR0    Y  | |t	|j
 qd S )NTr   r   r   z/Cannot broadcast a non-scalar to a scalar arrayr   rE  )rL  z=operands could not be broadcast together with remapped shapesr3  r(  r   rH  rJ   rJ   zDinput operand has more dimensions than allowed by the axis remappingrI  r]   z4all elements of broadcast shape must be non-negativer   )r]   rI   g?g@z7The second argument "shape" must be a tuple of integersr   z-The first argument "array" must be array-likerH   )rI   r3   z*object cannot be interpreted as an integer)r+   r   r&   r/  r   r   r   r   r   r   r   )	rr   r   r   r  r$   Ztarget_shapeerrr   r   r   r   r   test_broadcast_to_raises  sv    $*z.TestArrayManipulation.test_broadcast_to_raisesc                 C   s*   t dd }| }| }| || d S )Nc                   S   s   t ddS )Nr3   r   r(   r   r   r   r   _broadcast_to_1  s    zMTestArrayManipulation.test_broadcast_to_corner_cases.<locals>._broadcast_to_1)r   r  rz   )rr   rS  r|   r}   r   r   r   test_broadcast_to_corner_cases   s
    
z4TestArrayManipulation.test_broadcast_to_corner_casesc                 C   s\   t }tdd|}tjdtjd}d}|||}d|d< | | d | | d d S )	NTr   rI   r.  rI  r   r   r   )r+   r   r&   r/  r!  r   sum)rr   r   r   rJ  r*   viewr   r   r   test_broadcast_to_change_view	  s    
z3TestArrayManipulation.test_broadcast_to_change_viewc           	      C   s\   t }tdd|}tdddgg}|D ].\}}}||||}||||}| || q(d S )NTr   rI   rI  r   )r0   r   r&   r0  rz   )	rr   r   r   r  rJ  r*   r/   r|   r}   r   r   r   test_broadcast_to_indexing  s    z0TestArrayManipulation.test_broadcast_to_indexingc                 C   s8   t dd }td}||}||}| || d S )Nc                 S   s   t | d}||j|j|jfS )Nr   )r&   r)   r8   r*   rX   )r$   retr   r   r   r   "  s    z@TestArrayManipulation.test_broadcast_to_array_attrs.<locals>.foorH   )r   r&   rK   r  rz   )rr   r   r$   r|   r}   r   r   r   test_broadcast_to_array_attrs  s    


z3TestArrayManipulation.test_broadcast_to_array_attrsc                 C   s   t }tdd|}dgddgdgdgdgddgg dd	gd
dgddgddgddgddgddgddgddgddgddgddgddgddgddgddgddgddgddgddgg}|D ],}|| }|| }| |t | || qd S )NTr   r   r  r   rO  )rH   r   ))r   r  r   r   rJ   r[  )r   rJ   r  r\  rG  rH   rJ   r   r   r   r   r   rE  rI   r   )r-   r   ZassertIsInstancer   rz   )rr   r   r   r  input_shaper|   r}   r   r   r   test_broadcast_shapes-  sF    z+TestArrayManipulation.test_broadcast_shapesc              	   C   s   t }tdd|}|   ddgddgg ddd	gg d
ddgg}|D ]H}| t}||  W d    n1 sr0    Y  | dt|j qFd S )NTr   rE  r(  r   r3  rE  rE  r(  rJ   rH   r   rI   rH   rH   )r   r]  r   )r   r   rI   =shape mismatch: objects cannot be broadcast to a single shaper-   r   r   r   r   r   r   r   )rr   r   r   r  r^  r   r   r   r   test_broadcast_shapes_raisesW  s     &z2TestArrayManipulation.test_broadcast_shapes_raisesc                 C   sf   t }tdd|}|   | t}|ddd W d    n1 sF0    Y  | dt|j d S )NTr   r   rI   r   z#negative dimensions are not allowedrd  )rr   r   r   r   r   r   r   (test_broadcast_shapes_negative_dimensionn  s    *z>TestArrayManipulation.test_broadcast_shapes_negative_dimensionc              	   C   st   t }tdd|}|   g d}|D ]H}| t}||  W d    n1 sR0    Y  | dt|j q&d S )NTr   ))r   )r   )g333333@)stringr   r3   )rJ   rh  z#must be either an int or tuple[int])r-   r   r   r   r   r   r   r   )rr   r   r   Zinpsinpr   r   r   r   "test_broadcast_shapes_invalid_typex  s    &z8TestArrayManipulation.test_broadcast_shapes_invalid_typec                    s   t tdd  fdd}dD ]}t|}|| q&dddgddgd	d
ggddfD ]}|| q\t} d W d    n1 s0    Y  dt|j d S )NTr   c                    s    | } | } || d S r!   r   )r`   r|   r}   r9  r   r   r     s    z/TestArrayManipulation.test_shape.<locals>.check)r   r   r   )r   r   r   rJ   FrI   rH   r   r   rF  r3   z+The argument to np.shape must be array-like)	rj   r   r&   emptyr   r   r   r   r   )rr   r   tr$   r   r   r9  r   
test_shape  s    

"
&z TestArrayManipulation.test_shapec                    sJ   t }tdd|} fdd}| D ] }||}||} || q$d S )NTr   c                  3   s   t ddV  t jdddV  t g V   jd} d| t | dk< | V  | dddV  | jd	d
ddV  | dddd d
 V  | d V  d S )Nr   r   r   Z
fill_valued           皙?r   2   rI   r   r  r   rJ   r   )r&   rK   fullr   randomrandnabsr"   r2   rq   r   r   a_variations  s    zBTestArrayManipulation.test_flatnonzero_basic.<locals>.a_variations)rk   r   rz   rr   r   r   rw  r3   r|   r}   r   rq   r   test_flatnonzero_basic  s    
z,TestArrayManipulation.test_flatnonzero_basicc                    sJ   t }tdd|} fdd}| D ] }||}||} || q$d S )NTr   c                  3   s   t dddkV  t jdddV  t jdddV  t g V  t g dV   jd} | d	kV  | ddd
dkV  | jdddddkV  | dd
dd d dkV  | | d kV  | |  kV  d S )Nr   r   rI   r   rn  rJ   )g      rp  rb   ro  rq  r   g      ?rr  r   r  r   )r&   rK   rs  r   rt  ru  r"   r2   rq   r   r   rw    s    
z?TestArrayManipulation.test_argwhere_basic.<locals>.a_variations)rl   r   rz   rx  r   rq   r   test_argwhere_basic  s    
z)TestArrayManipulation.test_argwhere_basicc                   c   sf   dV  dV  ddgddgfV  dV  t jV  dV  dV  dV  d	V  d
V  dV  d V  dV  dV  d S )N)rP  )r@  g@)g      @gffffff@)rp  rb   rp  g      r   rJ   rI   rH   r   FTr;  y       @      ?Za_string )r&   r*  r   r   r   r   array_like_variations  s    z+TestArrayManipulation.array_like_variationsc                 C   s@   t }tdd|}|  D ] }||}||}| || qd S r   )rk   r   r|  rz   rr   r   r   r3   r|   r}   r   r   r   test_flatnonzero_array_like  s    z1TestArrayManipulation.test_flatnonzero_array_likec                 C   s@   t }tdd|}|  D ] }||}||}| || qd S r   )rl   r   r|  rz   r}  r   r   r   test_argwhere_array_like  s    z.TestArrayManipulation.test_argwhere_array_likec           	      C   sT   t }tdd|}dd |D }|| }|gt| }dd |D }| || d S )NTr   c                 S   s   g | ]}t |qS r   r&   r/  r   sr   r   r   r     r   zOTestArrayManipulation.broadcast_arrays_assert_correct_shape.<locals>.<listcomp>c                 S   s   g | ]
}|j qS r   r*   r   r3   r   r   r   r     r   r.   r   r   rz   )	rr   input_shapesr   r   r   inarrays	outarraysr|   r}   r   r   r   %broadcast_arrays_assert_correct_shape  s    z;TestArrayManipulation.broadcast_arrays_assert_correct_shapec                 C   sd   t }tdd|}g d}|D ]@}|g}| || ||g}| || |||g}| || qd S )NTr   )	r   rE  r   r   r   )rH   r   rG  r]  r?  )r.   r   r  )rr   r   r   r  r*   r  Zinput_shapes2Zinput_shapes3r   r   r   'test_broadcast_arrays_same_input_shapes   s    
z=TestArrayManipulation.test_broadcast_arrays_same_input_shapesc                 C   s   ddgdgddgdgddgdgddgdgddgdgddgdgddgdgddgdgd	dgdgdd	gdgddgdgddgdgdd	gd	gg}|D ]*\}}|  || |  |d d d
 | qd S )Nr   rE  rG  r?  r]  rO  r   r   r   r]   r  rr   r  r  r   r   r   r   9test_broadcast_arrays_two_compatible_by_ones_input_shapes  s"    












zOTestArrayManipulation.test_broadcast_arrays_two_compatible_by_ones_input_shapesc                 C   s   ddgdgddgdgddgdgddgdgddgdgddgdgddgdgddgdgddgdgddgdgddgdgdd	gd	gdd
gd	gdd	gd	gdd	gd	gddgdgdd
gd
gddgdgddgdgdd
gd
gg}|D ]*\}}|  || |  |d d d | qd S )Nr   rE  r?  r]  r   rO  rG  r   r   r   r   r]   r  r  r   r   r   Dtest_broadcast_arrays_two_compatible_by_prepending_ones_input_shapes8  s0    



















zZTestArrayManipulation.test_broadcast_arrays_two_compatible_by_prepending_ones_input_shapesc           	      C   sv   t }tdd|}ddgdgddgdgddgdgg}|D ]8\}}|| }d	d
 |D }|gt| }| || q8d S )NTr   Fr   rJ   rI   r   r3  c                 S   s   g | ]
}|j qS r   r  r  r   r   r   r   c  r   zLTestArrayManipulation.test_broadcast_arrays_scalar_input.<locals>.<listcomp>r  )	rr   r   r   r  r  r   r  r}   r|   r   r   r   "test_broadcast_arrays_scalar_inputY  s    


z8TestArrayManipulation.test_broadcast_arrays_scalar_inputc                 C   sB   t }tdd|}|dd}ddg}dd |D }| || d S )NTr   )rL  i  )i  r3  c                 S   s   g | ]
}|j qS r   r  r  r   r   r   r   l  r   zKTestArrayManipulation.test_broadcast_arrays_tuple_input.<locals>.<listcomp>)r.   r   rz   rr   r   r   r  r|   r}   r   r   r   !test_broadcast_arrays_tuple_inputg  s    
z7TestArrayManipulation.test_broadcast_arrays_tuple_inputc                 C   sT   t }tdd|}|tdtjdtjd}ddg}dd |D }| || d S )NTr   rI   rG  r.  c                 S   s   g | ]
}|j qS r   r  r  r   r   r   r   t  r   zOTestArrayManipulation.test_broadcast_arrays_non_array_input.<locals>.<listcomp>)r.   r   r&   Zintpr/  rz   r  r   r   r   %test_broadcast_arrays_non_array_inputo  s    z;TestArrayManipulation.test_broadcast_arrays_non_array_inputc                 C   sx   t }tdd|}|   | t.}tdd}d}||| W d    n1 sX0    Y  | dt	|j
 d S )NTr   r   r   zMismatch of argument types)r.   r   r   r   r   r&   rK   r"   r   r   r   )rr   r   r   r   r$   r7  r   r   r   /test_broadcast_arrays_invalid_mixed_input_typesw  s    (zETestArrayManipulation.test_broadcast_arrays_invalid_mixed_input_typesc                 C   sx   t }tdd|}|   | t.}tjdtjd}d}||| W d    n1 sX0    Y  | dt	|j
 d S )NTr   rH   r.  zhello worldzArgument "1" must be array-like)r.   r   r   r   r   r&   r/  rB  r   r   r   )rr   r   r   r   r$   r  r   r   r   #test_broadcast_arrays_invalid_input  s    (z9TestArrayManipulation.test_broadcast_arrays_invalid_inputc              
   C   s   t }tdd|}|   ddgddgg ddd	gg}|D ]n}||d d d
 fD ]V}| t&}dd |D }||  W d    n1 s0    Y  | dt|j qPq:d S )NTr   rE  r(  r   r3  r`  ra  rb  r]   c                 S   s   g | ]}t |qS r   r  r  r   r   r   r     r   zdTestArrayManipulation.test_broadcast_arrays_incompatible_shapes_raise_valueerror.<locals>.<listcomp>rc  )r.   r   r   r   r   r   r   r   )rr   r   r   r  r  r*   r   r  r   r   r   :test_broadcast_arrays_incompatible_shapes_raise_valueerror  s     &zPTestArrayManipulation.test_broadcast_arrays_incompatible_shapes_raise_valueerrorc                 C   sp   dd }|}t dd|}|tg dg dg dgd}|tg dg dg dgd}tj|| d S )	Nc                 S   sj   | j \}}t| |||f}t|dd}| }t|d|f}|d d|  ||d d d d f< |S Nr   rJ   r]   )r*   r&   r)   swapaxesr1   r"   r`   rc   rcr3   r7  ccdr   r   r   unfold_flatten  s    
$zITestArrayManipulation.test_readonly_after_flatten.<locals>.unfold_flattenTr   r&  rI   rI   rI   rH   rH   rH   rI   r   r&   r   r   Zassert_array_equal)rr   r  r   r   res_nbres_pyr   r   r   test_readonly_after_flatten  s    	""z1TestArrayManipulation.test_readonly_after_flattenc                 C   sp   dd }|}t dd|}|tg dg dg dgd}|tg dg dg dgd}tj|| d S )	Nc                 S   sj   | j \}}t| |||f}t|dd}| }t|d|f}|d d|  ||d d d d f< |S r  )r*   r&   r)   r  r6   r"   r  r   r   r   unfold_ravel  s    
$zETestArrayManipulation.test_readonly_after_ravel.<locals>.unfold_ravelTr   r&  r  r  rI   r  )rr   r  r   r   r  r  r   r   r   test_readonly_after_ravel  s    	""z/TestArrayManipulation.test_readonly_after_ravelc           	      C   s   |    tdd }|jdd}d|j_d|j_dd }|}tdd	|}||}||}d
}| |jj | d}| |jj| d S )Nr   r?  r   r  Fc                 S   s   |   }|S r!   r5   )r3   r   r   r   r   try_ravel_w_copy  s    zKTestArrayManipulation.test_mutability_after_ravel.<locals>.try_ravel_w_copyTr   zCNo copy was performed, so the resulting array must not be writeablez?A copy was performed, yet the resulting array is not modifiable)	r   r&   rK   r"   rD   r   r   r   
assertTrue)	rr   Za_cZa_fr  r   r   Zret_cZret_fr   r   r   r   test_mutability_after_ravel  s    z1TestArrayManipulation.test_mutability_after_ravel)A__name__
__module____qualname____doc__ro   r   r   r   r   r   r   r   r   r   r   r   r   enable_pyobj_flagsr   r  r  r  r  r  r  r  r  r  r  r#  r%  r<  rD  rK  rN  rR  rT  rW  rX  rZ  r_  re  rf  rj  rm  ry  rz  staticmethodr|  r~  r  r  r  r  r  r  r  r  r  r  r  r  r  r  __classcell__r   r   rs   r   rm      sx   p^/$D#,	H.#0	*

	!rm   __main__)N)F)B	functoolsr   	itertoolsr   numpyr&   ZunittestZnumba.core.compilerr   r   Znumbar   r   r   r	   Znumba.np.numpy_supportr
   Znumba.core.errorsr   Z
numba.corer   r   Znumba.tests.supportr   r   r   r   r  Zenable_pyobjectr   Znrtr    r%   r'   r+   r-   r.   r0   r4   r7   r9   r:   r=   r?   rC   rE   rF   rG   rP   rR   rS   rU   rV   rW   r\   r^   r_   rd   rf   ri   rj   rk   rl   rm   r  mainr   r   r   r   <module>   sv   	
          ;
