o
    0 i                     @   s   d dl Z d dl mZ d dlmZ d dlm  mZ d dlm	Z	 d dlm
Z
 d dlmZ d dlmZ dd	d
ededefddZdS )    N)Any)	DTypeLike)_array)_chunk)_data_transfer)_modesarrz_array.DistributedArrayaxisdtypereturnc                 C   s  d}t jt jt jt jd}|j|vrtd|j ||j }|t jt jfv r5| j|ur4| t j	} d}n| |} | j
}|rutt j|}|| j}| D ]$}	tt|	D ]}
t|	|
 jdkrdqX|	|
  |	|
< |	|
 | qXqP| jd | | j|d d   }d }i }| D ]\}}	g ||< |	D ]}| }|jd | |j|d d   }t|jtjr|jj}|d | ||d d   }tj||jj |}nt!"||j||d}|j}tj||# ||j$d}|| %| t|jdkr	 W d    q|jD ]8\}}|&|j' t!"||j||d}|j}t(j)||# |j$d}|d | ||d d   }|*|| qW d    n	1 sJw   Y  qqt+,||||| j-S )	NF)Zcupy_maxZcupy_minZcupy_sumZ	cupy_prodzUnsupported kernel: Tr      )r	   r
   )
prevent_gc).r   MAXZMINZSUMZPRODnameRuntimeError_modeZ_to_op_modeZREPLICAZ_chunks_maptypingcastZ_OpModeZidentity_ofr
   valuesrangelenZupdatescopyZ#set_identity_on_overwritten_entriesshapeitemsZon_readyindex
isinstancearrayr   Z_ArrayPlaceholderZ_ChunkZcreate_placeholderZdevice_manipulation_dimsZ
atleast_1drecordr   appendZ
wait_eventreadyr   Z
_AsyncDataZ
add_updater   ZDistributedArrayZ_comms)r   Zkernelr	   r
   Z
overwritesZmode_overridesmodeZ
chunks_mapidentitychunksir   Z	out_dtypeZout_chunks_mapdevchunkstreamZ	out_indexZ	old_shapeZ	out_shapeZ	out_chunkZ	out_arrayupdateZupdate_indexZout_update_arrayZ
out_update r)   n/home/app/PaddleOCR-VL-test/.venv_paddleocr/lib/python3.10/site-packages/cupyx/distributed/array/_reduction.py_execute   s   



 
 

%r+   )r   r   Znumpy.typingr   Zcupy._manipulation.dimsZ_manipulationdimsr   Zcupyx.distributed.arrayr   r   r   r   intr+   r)   r)   r)   r*   <module>   s"    