
    sgw                     >   d dl mZmZmZmZmZmZmZmZm	Z	 d dl
mZ d dlZd dlmZmZ g dZ ede      ded	efd
       Zded	efdZ edeee      Ze	 ddej2                  dee	eej6                  f      ded	eej2                  df   fd       Ze	 ddedee	eej6                  f      ded	ee   fd       ZddZ	 ddeedf   deeeef      dee	eej6                  f      ded	eeedf   eeeef   df   f   f
dZddede	eej6                  f   ded	efdZy)    )	AnyDictListOptionaloverloadSequenceTupleTypeVarUnion)
deprecatedN)GatherScatter)scatterscatter_kwargsgatherzC`is_namedtuple` is deprecated, please use the python checks instead)categoryobjreturnc                     t        |       S N)_is_namedtupler   s    S/var/www/html/venv/lib/python3.12/site-packages/torch/nn/parallel/scatter_gather.pyis_namedtupler      s     #    c                 Z    t        | t              xr t        | d      xr t        | d      S )N_asdict_fields)
isinstancetuplehasattrr   s    r   r   r      s+     	3V73	#:VwsI?Vr   T.inputstarget_gpusdimc                      y r    r#   r$   r%   s      r   r   r            r   c                      y r   r'   r(   s      r   r   r   )   r)   r   c                 >    fd	  |       }d|S # dw xY w)zSlice tensors into approximately equal chunks and distributes them across given GPUs.

    Duplicates references to objects that are not tensors.
    c                    t        | t        j                        rt        j                  d |       S t        |       r,t        t        |        D cg c]  } t        |       |  c}S t        | t              r)t        |       dkD  rt        t        t        |              S t        | t              r7t        |       dkD  r)t        t        |        D cg c]  }t        |       c}S t        | t              rKt        |       dkD  r=t        t        | j                                D cg c]  } t        |       |       c}S D cg c]  }|  c}S c c}w c c}w c c}w c c}w )Nr   )r   torchTensorr   applyr   zipmaptyper    lenlistdictitems)r   argsi_r%   scatter_mapr$   s       r   r:   zscatter.<locals>.scatter_map8   s   c5<<(==dC==#14c+s6K1LMIDIt$MMc5!c#hlSc2344c4 SX\%(#k3*?%@ADGAAc4 SX\*-s;		/L*MNQIDIaLNN())) N BN)s   EE+E 
	E%Nr'   )r#   r$   r%   resr:   s    `` @r   r   r   2   s*    *$&!J     kwargsc           	         | rt        | ||      ng }|rt        |||      ng }t        |      t        |      k  r7|j                  d t        t        |      t        |      z
        D               nMt        |      t        |       k  r6|j                  d t        t        |      t        |      z
        D               t	        |      t	        |      fS )z+Scatter with support for kwargs dictionary.c              3       K   | ]  }d   yw)r'   Nr'   .0r9   s     r   	<genexpr>z!scatter_kwargs.<locals>.<genexpr>[          
B 
   c              3       K   | ]  }i   y wr   r'   r@   s     r   rB   z!scatter_kwargs.<locals>.<genexpr>_   rC   rD   )r   r3   extendranger    )r#   r=   r$   r%   scattered_inputsscattered_kwargss         r   r   r   Q   s     =Cwv{C8<Bwv{C8
s#344  
c"23c:J6KKL 
 	
 
	V	,  
c"23c:J6KKL 
 	
 !"E*:$;;;r   outputstarget_devicec                 >    fd	  |       }d|S # dw xY w)a  Gather tensors from different GPUs on a specified device.

    This function is useful for gathering the results of a distributed computation.
    It takes a sequence of objects, one for each GPU, and returns a single object
    on the specified device.

    Args:
        outputs (Any): A sequence of objects (potentially tensors) to gather.
        target_device (Union[int, torch.device]): The device to gather the tensors to.
            Use 'cpu' for CPU to avoid a deprecation warning.
        dim (int, optional): The dimension along which to gather. Default: 0.

    Returns:
        Any: A gathered object (potentially tensor) on the specified device.
    c                      d   t        t        j                        rt        j                  g  S y t        t
              r:t        fd D              st        d       t               fdD              S t              r*t              j                  t        t                      S  t              t        t                      S )Nr   c              3   L   K   | ]  }t              t        |      k(    y wr   )r3   )rA   douts     r   rB   z-gather.<locals>.gather_map.<locals>.<genexpr>}   s     ;as3x3q6);s   !$z+All dicts must have the same number of keysc           	   3   ^   K   | ]  }| D cg c]  }||   	 c}      f ! y c c}w wr   r'   )rA   krO   
gather_maprJ   s      r   rB   z-gather.<locals>.gather_map.<locals>.<genexpr>   s-     S!a7,CaQqT,C!DES,Cs   -(-)r   r-   r.   r   r/   r5   all
ValueErrorr2   r   _maker1   r0   )rJ   rP   r%   rS   rK   s   `@r   rS   zgather.<locals>.gather_mapv   s    ajc5<<(<<s=W==;c4 ;7;; !NOO49SsSSS#9??3z3=#ABBtCyZg788r   Nr'   )rJ   rK   r%   r;   rS   s    `` @r   r   r   e   s*    "9 !
J 
r<   ).)r   ) typingr   r   r   r   r   r   r	   r
   r   typing_extensionsr   r-   torch.nn.parallel._functionsr   r   __all__FutureWarningboolr   r   r5   r4   r    r"   r.   intdevicer   strr   r   r'   r   r   <module>r`      s   W W W (  8 2 Is t 	
   CtU# 
 LL%U\\ 123 
 5<<	 
 
 %U\\ 123 
 
!W	 
F 	<#s(O<T#s(^$< %U\\ 123< 
	<
 5c?E$sCx.#"5667<(%C %c5<<.?(@ %s %SV %r   