§
    Ë²ÅjÓ  ã                   óÐ  — d dl mZ d dlmZmZmZmZmZmZm	Z	m
Z
mZmZ d dlZd dlmZ d dlmZ d dlmZ e
rd dlmZ d dlmZ d	gZd
edefd„Zd
edefd„Zdd„Zdd„Zdd
edee	e                  defd„Z	 d deej                 deee ej!        f                  dedeeej                          fd„Z" ede¬¦  «        Z#	 d de#deee ej!        f                  dedee#         fd„Z$dS )!é    )ÚOrderedDict)
ÚcastÚDictÚIteratorÚListÚOptionalÚSequenceÚSetÚTYPE_CHECKINGÚTypeVarÚUnionN©Ú_get_device_index)ÚModule)Úcomm)ÚScriptModule)ÚEnabledProxyÚ	replicateÚmoduleÚreturnc                 ó>   — dd l }t          | |j        j        ¦  «        S ©Nr   )Ú	torch.jitÚ
isinstanceÚjitr   ©r   Útorchs     úLC:\Projects\Historical\venv\Lib\site-packages\torch/nn/parallel/replicate.pyÚ_is_script_moduler      s#   € ØÐÐÐå�f˜eœiÔ4Ñ5Ô5Ð5ó    c                 ó>   — dd l }t          | |j        j        ¦  «        S r   )r   r   Ú_CÚScriptMethodr   s     r   Ú_is_script_methodr$   #   s#   € ØÐÐÐå�f˜eœhÔ3Ñ4Ô4Ð4r    r   c                  ó<   — dd l } | j                             ¦   «         S r   )r   r   r   ©r   s    r   Ú_init_script_moduler'   )   s"   € ØÐÐÐàŒ9×!Ò!Ñ#Ô#Ð#r    r   c                  ó,   — dd l } | j        j        j        S r   )Útorch.jit._stater   Ú_stateÚ_enabledr&   s    r   Ú_is_jit_enabledr,   /   s   € ØÐÐÐàŒ9ÔÔ$Ð$r    Úmemoc                 ó¦  — dt           dt          t                    fd„}t          ¦   «         sdS |€t          ¦   «         }|                     | ¦  «         t          | ¦  «        r@|                      || ¦  «        ¦  «         t          d„  || ¦  «        D ¦   «         ¦  «        S |                      ¦   «         D ]}||v rŒt          ||¦  «        s dS ŒdS )Nr   r   c                 óL   — |                       ¦   «         }t          |¦  «         |S ©N)ÚmodulesÚnext)r   Úgens     r   Údescendant_modulesz0_replicatable_module.<locals>.descendant_modules>   s!   € Ø�nŠnÑÔˆÝˆS‰	Œ	ˆ	Øˆ
r    Tc              3   ó4   K  — | ]}t          |¦  «        V — Œd S r0   )r   )Ú.0Ú
descendants     r   ú	<genexpr>z'_replicatable_module.<locals>.<genexpr>L   s<   è è € ð 
ð 
Ø.8Õ˜jÑ)Ô)ð
ð 
ð 
ð 
ð 
ð 
r    F)
r   r   r,   ÚsetÚaddr   ÚupdateÚallÚchildrenÚ_replicatable_module)r   r-   r4   Úchilds       r   r>   r>   <   s  € ð¥6ð ­hµvÔ.>ð ð ð ð õ
 ÑÔð ØˆtØ€|Ý‰uŒuˆð 	‡H‚HˆVÑÔÐÝ˜Ñ Ô ð 
Ø�ŠÐ&Ð& vÑ.Ô.Ñ/Ô/Ð/Ýð 
ð 
Ø<NÐ<NÈvÑ<VÔ<Vð
ñ 
ô 
ñ 
ô 
ð 	
ð —’Ñ"Ô"ð ð ˆð �Dˆ=ˆ=ØÝ# E¨4Ñ0Ô0ð 	Ø�5�5ð	ð ˆ4r    FÚtensorsÚdevicesÚdetachc                 óö   ‡ ‡— ddl m} |rt          j        ‰ |¦  «        S t	          ‰ ¦  «        dk    rF |j        |g‰ ¢R Ž Šˆˆ fd„t          dt	          ‰¦  «        t	          ‰ ¦  «        ¦  «        D ¦   «         S g S )Nr   )Ú	Broadcastc                 óD   •— g | ]}‰||t          ‰¦  «        z   …         ‘ŒS © )Úlen)r6   ÚiÚtensor_copiesr@   s     €€r   ú
<listcomp>z0_broadcast_coalesced_reshape.<locals>.<listcomp>h   s>   ø€ ð ð ð àð ˜a !¥c¨'¡l¤lÑ"2Ð2Ô3ðð ð r    )Útorch.nn.parallel._functionsrD   r   Úbroadcast_coalescedrG   ÚapplyÚrange)r@   rA   rB   rD   rI   s   `   @r   Ú_broadcast_coalesced_reshaperO   [   s°   øø€ ð
 7Ð6Ð6Ð6Ð6Ð6àð ÝÔ'¨°Ñ9Ô9Ð9õ ˆw‰<Œ<˜!ÒÐØ+˜IœO¨GÐ>°gÐ>Ð>Ð>ˆMðð ð ð ð å˜q¥# mÑ"4Ô"4µc¸'±l´lÑCÔCðñ ô ð ð
 ˆIr    ÚT)ÚboundÚnetworkc           	      óÄ  ‡— t          | ¦  «        st          d¦  «        ‚|sg S d„ |D ¦   «         }t          |¦  «        }t          |                      ¦   «         ¦  «        }d„ t          |¦  «        D ¦   «         }t          |||¦  «        }t          |                      ¦   «         ¦  «        }g }g }	|D ]6}
|
j        r|s| 	                    |
¦  «         Œ!|	 	                    |
¦  «         Œ7d„ t          |¦  «        D ¦   «         }d„ t          |	¦  «        D ¦   «         }t          |||¬¦  «        }t          |	|d¬¦  «        }t          |  
                    ¦   «         ¦  «        }d„ |D ¦   «         Ši }t          |¦  «        D ]^\  }}|||<   t          |¦  «        D ]D}|                     ¦   «         }t          ¦   «         |_        ‰|          	                    |¦  «         ŒEŒ_t          |¦  «        D �]Û\  }}|j                             ¦   «         D ]w\  }}|€+t          |¦  «        D ]}‰|         |         }d |j        |<   ŒŒ2||         }t          |¦  «        D ]-}‰|         |         }t#          ||‰|         |         ¦  «         Œ.Œx|j                             ¦   «         D ]ƒ\  }}|€+t          |¦  «        D ]}‰|         |         }d |j        |<   ŒŒ2||         }t          |¦  «        D ]9}‰|         |         }||         |         }t#          |||¦  «         ||j        |<   Œ:Œ„|j                             ¦   «         D ]�\  }}
|
€+t          |¦  «        D ]}‰|         |         }d |j        |<   ŒŒ2|
j        r|s|}||
         }n
|}||
         }t          |¦  «        D ]-}‰|         |         }t#          ||||         |         ¦  «         Œ.ŒŽ�ŒÝˆfd	„t          |¦  «        D ¦   «         S )
NzKCannot replicate network where python modules are childrens of ScriptModulec                 ó.   — g | ]}t          |d ¦  «        ‘ŒS )Tr   )r6   Úxs     r   rJ   zreplicate.<locals>.<listcomp>�   s#   € Ð;Ð;Ð;¨aÕ   DÑ)Ô)Ð;Ð;Ð;r    c                 ó   — i | ]\  }}||“Œ	S rF   rF   )r6   ÚidxÚparams      r   ú
<dictcomp>zreplicate.<locals>.<dictcomp>…   s   € ÐDÐDÐD¡J C¨�U˜CÐDÐDÐDr    c                 ó   — i | ]\  }}||“Œ	S rF   rF   ©r6   rW   Úbufs      r   rY   zreplicate.<locals>.<dictcomp>‘   s   € ÐHÐHÐH¡h c¨3˜˜cÐHÐHÐHr    c                 ó   — i | ]\  }}||“Œ	S rF   rF   r[   s      r   rY   zreplicate.<locals>.<dictcomp>’   s   € ÐPÐPÐP©(¨#¨s˜S #ÐPÐPÐPr    )rB   Tc                 ó   — g | ]}g ‘ŒS rF   rF   )r6   Ú_s     r   rJ   zreplicate.<locals>.<listcomp>š   s   € Ð(=Ð(=Ð(=°¨Ð(=Ð(=Ð(=r    c                 óR   •— g | ]#}t          t          ‰|         d          ¦  «        ‘Œ$S )r   )r   rP   )r6   ÚjÚmodule_copiess     €r   rJ   zreplicate.<locals>.<listcomp>Ô   s-   ø€ ÐFÐFÐF¨Q�D•�M !Ô$ QÔ'Ñ(Ô(ÐFÐFÐFr    )r>   ÚRuntimeErrorrG   ÚlistÚ
parametersÚ	enumeraterO   ÚbuffersÚrequires_gradÚappendr1   rN   Ú_replicate_for_data_parallelr   Ú_former_parametersÚ_modulesÚitemsÚsetattrÚ_parametersÚ_buffers)rR   rA   rB   Únum_replicasÚparamsÚparam_indicesÚparam_copiesrg   Ú
buffers_rgÚbuffers_not_rgr\   Úbuffer_indices_rgÚbuffer_indices_not_rgÚbuffer_copies_rgÚbuffer_copies_not_rgr1   Úmodule_indicesrH   r   ra   ÚreplicaÚkeyr?   Ú
module_idxrX   Ú	param_idxÚ
param_copyÚbuffer_copiesÚ
buffer_idxrb   s                                @r   r   r   s   sã  ø€ õ
   Ñ(Ô(ð 
Ýð(ñ
ô 
ð 	
ð
 ð Øˆ	à;Ð;°7Ð;Ñ;Ô;€GÝ�w‘<”<€Lå�'×$Ò$Ñ&Ô&Ñ'Ô'€FØDÐDµ)¸FÑ2CÔ2CÐDÑDÔD€MÝ/°¸ÀÑHÔH€Lå�7—?’?Ñ$Ô$Ñ%Ô%€GØ%'€JØ)+€NØð 'ð 'ˆØÔð 	' Vð 	'Ø×Ò˜cÑ"Ô"Ð"Ð"à×!Ò! #Ñ&Ô&Ð&Ð&àHÐHµ)¸JÑ2GÔ2GÐHÑHÔHÐØPÐPµiÀÑ6OÔ6OÐPÑPÔPÐå3°JÀÐPVÐWÑWÔWÐÝ7Ø˜¨ðñ ô Ðõ �7—?’?Ñ$Ô$Ñ%Ô%€GØ(=Ð(=°WÐ(=Ñ(=Ô(=€MØ(*€Nå˜wÑ'Ô'ð -ð -‰	ˆˆ6Ø!"ˆ�vÑÝ�|Ñ$Ô$ð 		-ð 		-ˆAØ×9Ò9Ñ;Ô;ˆGõ *5©¬ˆGÔ&à˜!Ô×#Ò# GÑ,Ô,Ð,Ð,ð		-õ ˜wÑ'Ô'ð (Hñ (H‰	ˆˆ6Ø œ/×/Ò/Ñ1Ô1ð 		Hð 		H‰JˆC�Øˆ}Ý˜|Ñ,Ô,ð 1ð 1�AØ+¨AÔ.¨qÔ1�GØ,0�GÔ$ SÑ)Ð)ð1ð ,¨EÔ2�
Ý˜|Ñ,Ô,ð Hð H�AØ+¨AÔ.¨qÔ1�GÝ˜G S¨-¸Ô*:¸:Ô*FÑGÔGÐGÐGðHð !Ô,×2Ò2Ñ4Ô4ð 	Að 	A‰JˆC�Øˆ}Ý˜|Ñ,Ô,ð 4ð 4�AØ+¨AÔ.¨qÔ1�GØ/3�GÔ'¨Ñ,Ð,ð4ð *¨%Ô0�	Ý˜|Ñ,Ô,ð Að A�AØ+¨AÔ.¨qÔ1�GØ!-¨a¤°Ô!;�Jõ ˜G S¨*Ñ5Ô5Ð5à6@�GÔ.¨sÑ3Ð3ðAð œ×-Ò-Ñ/Ô/ð 	Hð 	H‰HˆC�Øˆ{Ý˜|Ñ,Ô,ð 1ð 1�AØ+¨AÔ.¨qÔ1�GØ,0�GÔ$ SÑ)Ð)ð1ð Ô$ð <¨Vð <Ø$4�MØ!2°3Ô!7�J�Jà$8�MØ!6°sÔ!;�JÝ˜|Ñ,Ô,ð Hð H�AØ+¨AÔ.¨qÔ1�GÝ˜G S¨-¸Ô*:¸:Ô*FÑGÔGÐGÐGðHñ	Hð  GÐFÐFÐFµ%¸Ñ2EÔ2EÐFÑFÔFÐFr    )r   r   )r   r   r0   )F)%Úcollectionsr   Útypingr   r   r   r   r   r	   r
   r   r   r   r   Útorch._utilsr   Útorch.nn.modulesr   Útorch.nn.parallelr   r   r   r)   r   Ú__all__Úboolr   r$   r'   r,   r>   ÚTensorÚintÚdevicerO   rP   r   rF   r    r   ú<module>r�      sŸ  ðØ #Ð #Ð #Ð #Ð #Ð #ðð ð ð ð ð ð ð ð ð ð ð ð ð ð ð ð ð ð ð ð ð ð ð ð €€€Ø *Ð *Ð *Ð *Ð *Ð *Ø #Ð #Ð #Ð #Ð #Ð #Ø "Ð "Ð "Ð "Ð "Ð "ð ð .Ø&Ð&Ð&Ð&Ð&Ð&Ø-Ð-Ð-Ð-Ð-Ð-ð ˆ-€ð6˜fð 6¨ð 6ð 6ð 6ð 6ð5˜fð 5¨ð 5ð 5ð 5ð 5ð$ð $ð $ð $ð%ð %ð %ð %ðð  ð ¨x¸¸F¼Ô/Dð ÐPTð ð ð ð ðD ðð Ø�e”lÔ#ðà�e˜C ¤Ð-Ô.Ô/ðð ðð 
ˆ$ˆuŒ|Ô
Ôð	ð ð ð ð* €GˆC�vÐÑÔ€ð ðaGð aGØðaGà�e˜C ¤Ð-Ô.Ô/ðaGð ðaGð 
ˆ!„Wð	aGð aGð aGð aGð aGð aGr    