o
    ۾is                     @   sX   d dl Zd dlmZ d dlmZ d dlmZmZ G dd deZ	e
dkr*e  dS dS )    N)
namedtuple)cuda)unittestCUDATestCasec                   @   sd   e Z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S )TestCudaArrayArgc                    sn   t jdddddd  t d fdd}tjd	tjd
}t|}|d || | t||k d S )Nzdouble(double[:],int64)T)deviceinlinec                 S   s   | | S )N )acr	   r	   [/home/ubuntu/.local/lib/python3.10/site-packages/numba/cuda/tests/cudapy/test_array_args.pydevice_function   s   z8TestCudaArrayArg.test_array_ary.<locals>.device_functionzvoid(double[:],double[:])c                    s   t d} | |||< d S )N   )r   grid)xyir   r	   r   kernel   s   
z/TestCudaArrayArg.test_array_ary.<locals>.kernel
   dtype)r   r   )r   jitnparangedouble
zeros_like
assertTrueall)selfr   r   r   r	   r   r   test_array_ary	   s   

zTestCudaArrayArg.test_array_aryc                 C   s^   t jdd }d}tjt|tjd}|d || tt|D ]}| || ||  q d S )Nc                 S   s(   |d | d< |d | d< |d | d< d S )Nr   r      r	   rr   r	   r	   r   f   s   z)TestCudaArrayArg.test_unituple.<locals>.f)r   r!      r   r   r   )r   r   r   zeroslenint64rangeassertEqual)r   r$   r   r#   r   r	   r	   r   test_unituple   s   
zTestCudaArrayArg.test_unituplec                 C   s   t jdd }d}tjt|d tjd}tjt|d tjd}|d ||| tt|D ]}| || ||  q/tt|D ]}| || ||t|   qBd S )Nc                 S   sL   |d | d< |d | d< |d | d< |d |d< |d |d< |d |d< d S Nr   r   r!   r%         r	   r1r2r   r	   r	   r   r$   (   s   z&TestCudaArrayArg.test_tuple.<locals>.f)r   r!   r%   g      @g      @g      @r!   r   r&   )	r   r   r   r'   r(   r)   float64r*   r+   )r   r$   r   r1   r2   r   r	   r	   r   
test_tuple'   s   
zTestCudaArrayArg.test_tuplec                 C   sl   t jdd }tdd}|dd}tjt|tjd}|d || | |d	 |j | |d |j	 d S )
Nc                 S   s   |j | d< |j| d< d S Nr   r   r   r   r"   r	   r	   r   r$   =   s   
z.TestCudaArrayArg.test_namedunituple.<locals>.fPointr6   r   r!   r   r&   r   )
r   r   r   r   r'   r(   r)   r+   r   r   )r   r$   r7   r   r#   r	   r	   r   test_namedunituple<   s   


z#TestCudaArrayArg.test_namedunituplec                 C   s   t jdd }tdd}|ddd}tjdtjd}tjdtjd}|d	 ||| | |d
 |j | |d |j	 | |d
 |j
 d S )Nc                 S   s"   |j | d< |j| d< |j|d< d S r5   r   r   r#   r0   r	   r	   r   r$   K   s   

z+TestCudaArrayArg.test_namedtuple.<locals>.fr7   r9   r   r!   gS@r   r&   r   )r   r   r   r   r'   r)   r3   r+   r   r   r#   )r   r$   r7   r   r1   r2   r	   r	   r   test_namedtupleJ   s   

z TestCudaArrayArg.test_namedtuplec                 C   sF   t jdd }t }tjdtjd}|d || | |d d d S )Nc                 S   s   t || d< d S )Nr   r(   r"   r	   r	   r   r$   \   s   z,TestCudaArrayArg.test_empty_tuple.<locals>.fr   r   r&   r   )r   r   tupler   onesr)   r+   r   r$   r   r#   r	   r	   r   test_empty_tuple[   s   
z!TestCudaArrayArg.test_empty_tuplec                 C   sT   t jdd }d}tjdtjd}|d || | |d d | |d	 d d S )
Nc                 S   s    t || d< t |d | d< d S r5   r;   r"   r	   r	   r   r$   g   s   z6TestCudaArrayArg.test_tuple_of_empty_tuples.<locals>.f)r	   r	   r	   r!   r   r&   r   r%   r   r   r   r   r=   r)   r+   r>   r	   r	   r   test_tuple_of_empty_tuplesf   s   
z+TestCudaArrayArg.test_tuple_of_empty_tuplesc                 C   s   t jdd }d}tjdtjd}|d || | |d d | |d	 d | |d
 d
 | |d d | |d d | |d d | |d d | |d d | |d d d S )Nc                 S   s   t || d< t |d | d< t |d | d< t |d | d< |d d | d< |d d | d< |d d | d< |d d | d< |d d | d	< d S )
Nr   r   r!   r%   r.   r/            r;   r"   r	   r	   r   r$   t   s   z0TestCudaArrayArg.test_tuple_of_tuples.<locals>.f)r	   )r/   rB   )rD   	   r   rE   r   r&   r   r%   r   r!   r.   r/   rB   rD   rC   r   r@   r>   r	   r	   r   test_tuple_of_tupless   s   
z%TestCudaArrayArg.test_tuple_of_tuplesc                 C   s   t jdd }d}tjdtjd}|d || | |d d | |d	 d
 | |d d | |d
 d | |d d | |d d d S )Nc                 S   s\   t || d< t |d | d< |d d | d< |d d | d< |d d | d< |d | d< d S r-   r;   r"   r	   r	   r   r$      s   z<TestCudaArrayArg.test_tuple_of_tuples_and_scalars.<locals>.f))rB   r/   r.   rC   rE   r   r&   r   r!   r   r%   rB   r/   r.   rC   r@   r>   r	   r	   r    test_tuple_of_tuples_and_scalars   s   
z1TestCudaArrayArg.test_tuple_of_tuples_and_scalarsc                 C   s^   t jdd }d}t|}t|}|d }|||f}|d|f | tj|||  d S )Nc                 S   sB   t d}|t| d k r| d | | d |  | d |< d S d S )Nr   r   r!   )r   r   r(   )r   r   r	   r	   r   r$      s   
$z0TestCudaArrayArg.test_tuple_of_arrays.<locals>.fr   r%   r   )r   r   r   r'   	ones_liketestingassert_equal)r   r$   Nx0x1x2r   r	   r	   r   test_tuple_of_arrays   s   



z%TestCudaArrayArg.test_tuple_of_arraysc                 C   s   t jdd }tjdtjd}d| ddf}tjdtjd}|d || | |d	 d	 | |d
 d | |d d | |d d | |d d d S )Nc                 S   sP   |d d | d< |d d | d< |d | d< |d d | d< |d d | d< d S )Nr   r   r!   r%   r.   r	   r"   r	   r	   r   r$      s
   z<TestCudaArrayArg.test_tuple_of_array_scalar_tuple.<locals>.fr!   r   r   )r.   r%   r/   r&   r   r   r%   r.   )r   r   r   r   r)   r'   r+   )r   r$   zr   r#   r	   r	   r    test_tuple_of_array_scalar_tuple   s   
z1TestCudaArrayArg.test_tuple_of_array_scalar_tupleN)__name__
__module____qualname__r    r,   r4   r8   r:   r?   rA   rF   rG   rO   rQ   r	   r	   r	   r   r      s    r   __main__)numpyr   collectionsr   numbar   numba.cuda.testingr   r   r   rR   mainr	   r	   r	   r   <module>   s     A