Ë
    L²jöI  ã                   óJ  — d dl mZ d dlZd dlmc mZ d dlmZ d dlm	Z	 d dl
mZ d dlmZ d dlmZmZ d dlmZ d d	lmZmZmZ d d
lmZmZ d„ Zdd„Zdej8                  fd„Zdej8                  fd„Zd„ Zd„ Z d„ Z!d„ Z" e!«       Z# e"«       Z$d„ Z% e%«       Z&d„ Z' e'«       Z(d„ Z)y)é    )ÚcastN)Ú_prims)ÚDispatchKey)Úautograd_not_implemented)ÚHigherOrderOperator)ÚCUDARngStateHelperÚmake_contiguous_strides_for)ÚFakeTensorMode)Údisable_proxy_modes_tracingÚProxyTorchDispatchModeÚtrack_tensor_tree)Ú_deviceÚ_dtypec           	      óh   — t        d| j                  › d| j                  › d| j                  › d�«      ‚)Nz"You are trying to functionalize a z RNG operator but zE does not use Philox/counter-based RNG. Therefore, functionalizing a zo RNG operator is not supported. We are discussing the possibility of a Philox-based RNG implementation for CPU.)ÚRuntimeErrorÚtype)Údevices    úU/var/www/html/ia-enhanced/venv/lib/python3.12/site-packages/torch/_prims/rng_prims.pyÚthrow_on_non_cudar      sB   € Ü
Ø
,¨V¯[©[¨MÐ9KÈFÏKÉKÈ=ð YFØFLÇkÁkÀ]ð Sið	ióð ó    c                 óš  — t         j                  j                  d| z   |d|¬«      }|j                  |«       t	        t         j
                  j                  j                  | «      }|j                  }|r||_	        ||fD ]J  }	||	_
        t         j                  j                  j                  |	_        | |z   |	_        ||	_        ||	_        ŒL y )Nz
rngprims::© )Úmutates_argsÚschema)ÚtorchÚlibraryÚ	custom_opÚregister_fakeÚgetattrÚ_opsÚopsÚrngprimsÚdefaultÚ_tagsÚ__doc__Ú_prims_commonÚRETURN_TYPEÚNEWÚreturn_typer   Ú	impl_atenÚprim_meta_impl)
Únamer   r*   Ú	impl_metaÚdocÚtagsÚrngprim_defÚprim_packetÚprimÚps
             r   Úregister_rng_primr4      s¶   € Ü—-‘-×)Ñ)Ø�tÑ˜Y°RÀð *ó €Kð ×Ñ˜iÔ(äœ%Ÿ*™*Ÿ.™.×1Ñ1°4Ó8€KØ×Ñ€DÙØˆŒ
à˜4Ð ò %ˆØˆŒ	Ü×+Ñ+×7Ñ7×;Ñ;ˆŒà˜&‘=ˆŒØˆŒØ$ˆÕñ%r   Úshapec                 ór   — t        j                  t        j                  dt        j                  ¬«      «      S )Nr   ©Údtype)r   Ú
TensorLiker   ÚtensorÚint64)r5   s    r   Úphilox_rand_offset_metar<   3   s$   € ô ×ÑœUŸ\™\¨!´5·;±;Ô?Ó@Ð@r   c                 ó˜  — d}| D ]  }||z  }Œ	 t        j                  |t         j                  ¬«      }d}d}d}t         j                  j	                  t         j                  j                  «       «      }|j                  |z  }t        t        |«      }	|	|z   dz
  |z  }
t        |
|j                  |z  «      }
|	dz
  ||
z  |z  z  dz   |z  S )Né   r7   é   é   )r   Úscalar_tensorr;   ÚcudaÚget_device_propertiesÚcurrent_deviceÚmax_threads_per_multi_processorr   ÚintÚminÚmulti_processor_count)r5   Únumel_scalarÚdim_sizeÚnumelÚ
block_sizeÚunrollÚcurand4_engine_callsÚdevice_propertyÚblocks_per_smÚnumÚ	grid_sizes              r   Úphilox_rand_offsetrS   9   sÔ   € ð €LØò !ˆØ˜Ñ ‰ð!ä×Ñ ´E·K±KÔ@€Eà€JØ€FØÐÜ—j‘j×6Ñ6´u·z±z×7PÑ7PÓ7RÓS€OØ#×CÑCÀzÑQ€MÜ
Œs�EÓ
€CØ�zÑ! AÑ%¨*Ñ4€IÜ�I˜×DÑDÀ}ÑTÓU€IØ�1‰W˜* yÑ0°6Ñ9Ñ:¸QÑ>ÐBVÑVÐVr   c                  ó¨  — d} d}dt         j                  dt         j                  dt         j                  dt        t        df   d z  dt
        d	t        fd
„}dt         j                  dt         j                  dt         j                  dt        t        df   d z  dt
        d	t        fd„}t        | |||dt         j                  j                  f¬«       y )NÚphilox_randz{(SymInt[] size, Tensor seed, Tensor offset, int[]? stride, Device? device=None, ScalarType? dtype=None) -> (Tensor, Tensor)r5   ÚseedÚoffsetÚstride.r   r8   c                 óˆ   — |�t        d|› �«      ‚t        | «      }t        j                  | |||¬«      }t	        | «      }||fS )Nústride must be None, got )r5   Ústridesr8   r   )ÚAssertionErrorr	   r   Ú
TensorMetar<   )r5   rV   rW   rX   r   r8   Úrandom_valuess          r   Ú_philox_rand_metaz/register_philox_rand.<locals>._philox_rand_metaS   sW   € ð ÐÜ Ð#<¸V¸HÐ!EÓFÐFÜ,¨UÓ3ˆÜ×)Ñ)Ø ¨u¸Vô
ˆô )¨Ó/ˆØ˜vÐ&Ð&r   c                 ó^  — |�t        d|› �«      ‚|j                  dk(  rg }n|g}|j                  dk7  rt        |«      ‚t        j                  j                  |«      5  t        j                  ||«       t        j                  | ||¬«      }d d d «       t        | «      fS # 1 sw Y   ŒxY w)NrZ   ÚcpurB   )r   r8   )
r\   r   r   r   ÚrandomÚfork_rngr   Úset_torch_state_tensorÚrandrS   )r5   rV   rW   rX   r   r8   Údevicesr^   s           r   Ú_philox_randz*register_philox_rand.<locals>._philox_rande   s¬   € ð ÐÜ Ð#<¸V¸HÐ!EÓFÐFØ�;‰;˜%ÒØ‰Gà�hˆGà�;‰;˜&Ò Ü# FÓ+Ð+ä�\‰\×"Ñ" 7Ó+ñ 	JÜ×5Ñ5°d¸FÔCÜ!ŸJ™J u°VÀ5ÔIˆM÷	Jð Ô0°Ó7Ð7Ð7÷		Jð 	Jús   Á/B#Â#B,z$Philox based stateless rand operator)r,   r   r*   r-   r.   r/   )
r   ÚSizeÚTensorÚtuplerF   r   r   r4   ÚTagÚnondeterministic_seeded)r,   r   r_   rg   s       r   Úregister_philox_randrm   O   sä   € Ø€Dð K€Fð'Ü�z‰zð'ä�l‰lð'ô —‘ð'ô ”c˜3�h‘ $Ñ&ð	'ô
 ð'ô ó'ð$8Ü�z‰zð8ä�l‰lð8ô —‘ð8ô ”c˜3�h‘ $Ñ&ð	8ô
 ð8ô ó8ô2 ØØØØ#Ø2Ü�i‰i×/Ñ/Ð1ör   c                 óÈ  — |j                  d«      rB|j                  d«      }t        |t        «      rt        j                  |«      }|j
                  S | D �ch c]3  }t        |t        j                  «      sŒ|j                  j
                  ’Œ5 }}t        d„ |D «       «      ryt        d„ |D «       «      ryt        d„ |D «       «      ryt        d„ |D «       «      ry	y c c}w )
Nr   c              3   ó&   K  — | ]	  }|d k(  –— Œ y­w)rB   Nr   ©Ú.0Údevs     r   ú	<genexpr>zget_device.<locals>.<genexpr>�   s   è ø€ Ò
,˜Sˆ3�&�=Ñ
,ùó   ‚rB   c              3   ó&   K  — | ]	  }|d k(  –— Œ y­w)ÚxpuNr   rp   s     r   rs   zget_device.<locals>.<genexpr>’   ó   è ø€ Ò-˜cˆS�E�\Ñ-ùrt   rv   c              3   ó&   K  — | ]	  }|d k(  –— Œ y­w)ÚhpuNr   rp   s     r   rs   zget_device.<locals>.<genexpr>”   rw   rt   ry   c              3   ó&   K  — | ]	  }|d k(  –— Œ y­w)ra   Nr   rp   s     r   rs   zget_device.<locals>.<genexpr>–   rw   rt   ra   )ÚgetÚ
isinstanceÚstrr   r   r   ri   Úany)ÚargsÚkwargsr   Úargrf   s        r   Ú
get_devicer‚   ˆ   s³   € Ø‡z�z�(ÔØ—‘˜HÓ%ˆÜ�fœcÔ"Ü—\‘\ &Ó)ˆFØ�{‰{Ðà*.ÖP 3´*¸SÄ%Ç,Á,Õ2Oˆs�z‰z�‹ÐP€GÐPÜ
Ñ
, GÔ
,Ô,ØÜ	Ñ- WÔ-Ô	-ØÜ	Ñ- WÔ-Ô	-ØÜ	Ñ- WÔ-Ô	-ØØùò Qs   ÁCÁ8Cc                  ó�  ‡‡‡‡‡‡—  G d„ dt         «      }  | «       Š ‰j                  t        j                  «      t	        ‰d¬«      «       ‰j                  t        j
                  «      d„ «       Š‰j                  t        j                  «      d„ «       Š‰j                  t        j                  «      d„ «       Š‰j                  t        j                  «      d„ «       Š‰j                  t        j                  «      ˆˆˆˆfd	„«       Š‰j                  t        «      ˆfd
„«       }‰j                  t        «      ˆˆfd„«       }‰S )Nc                   ó(   ‡ — e Zd Zˆ fd„Zˆ fd„Zˆ xZS )ú>register_run_and_save_rng_state_op.<locals>.RunAndSaveRngStatec                 ó(   •— t         ‰| �  dd¬«       y )NÚrun_and_save_rng_stateT©Ú	cacheable©ÚsuperÚ__init__©ÚselfÚ	__class__s    €r   rŒ   zGregister_run_and_save_rng_state_op.<locals>.RunAndSaveRngState.__init__�   s   ø€ Ü‰GÑÐ5ÀÐÕFr   c                 ó*   •— t        ‰| �  |g|¢­i |¤ŽS ©N©r‹   Ú__call__)rŽ   Úopr   r€   r�   s       €r   r“   zGregister_run_and_save_rng_state_op.<locals>.RunAndSaveRngState.__call__    s   ø€ ä‘7Ñ# BÐ8¨Ò8°Ñ8Ð8r   ©Ú__name__Ú
__module__Ú__qualname__rŒ   r“   Ú__classcell__©r�   s   @r   ÚRunAndSaveRngStater…   œ   s   ø„ ô	G÷	9ð 	9r   r›   T©Údeferred_errorc                 óN   — t         j                  j                  «        | |i |¤ŽfS r‘   )r   rB   Úget_rng_state©r”   r   r€   s      r   Ú	impl_cudaz5register_run_and_save_rng_state_op.<locals>.impl_cudaª   s$   € ä�z‰z×'Ñ'Ó)©2¨tÐ+>°vÑ+>Ð>Ð>r   c                 ó:   — t        j                  «        | |i |¤ŽfS r‘   )r   rŸ   r    s      r   Úimpl_cpuz4register_run_and_save_rng_state_op.<locals>.impl_cpu®   s    € ä×"Ñ"Ó$¡b¨$Ð&9°&Ñ&9Ð9Ð9r   c                 ó„   — t        t        d«      r&t        j                  j                  «        | |i |¤ŽfS t	        d«      ‚©Nry   z2functionalize a hpu RNG operator is not supported.)Úhasattrr   ry   rŸ   r   r    s      r   Úimpl_hpuz4register_run_and_save_rng_state_op.<locals>.impl_hpu²   s:   € ä”5˜%Ô Ü—9‘9×*Ñ*Ó,©b°$Ð.A¸&Ñ.AÐAÐAÜÐOÓPÐPr   c                 óN   — t         j                  j                  «        | |i |¤ŽfS r‘   )r   rv   rŸ   r    s      r   Úimpl_xpuz4register_run_and_save_rng_state_op.<locals>.impl_xpu¸   s$   € ä�y‰y×&Ñ&Ó(©"¨dÐ*=°fÑ*=Ð=Ð=r   c                 óp   •— ‰‰‰‰	dœ}t        ||«      }||vrt        d|› �«      ‚||   } || g|¢­i |¤ŽS ©N)rB   ra   ry   rv   zBackend not supported for ©r‚   r\   )
r”   r   r€   Úimpl_mapr   Úimplr£   r¡   r§   r©   s
         €€€€r   Úimpl_backend_selectz?register_run_and_save_rng_state_op.<locals>.impl_backend_select¼   s`   ø€ ð ØØØñ	
ˆô ˜D &Ó)ˆØ˜Ñ!Ü Ð#=¸f¸XÐ!FÓGÐGØ˜ÑˆÙ�BÐ(˜Ò( Ñ(Ð(r   c                 óL   •— | 5   ‰|g|¢­i |¤Žcd d d «       S # 1 sw Y   y xY wr‘   r   )Úmoder”   r   r€   r¯   s       €r   Úimpl_fake_tensor_modezAregister_run_and_save_rng_state_op.<locals>.impl_fake_tensor_modeÊ   s/   ø€ ð ñ 	<Ù& rÐ;¨DÒ;°FÑ;÷	<÷ 	<ò 	<ús   „š#c                 ó:  •—  ‰|g|¢­i |¤Ž}t        j                  | j                  j                  |g|¢­«      }t        j                  | j                  j                  |«      }| j                  j	                  d‰	||«      }t        ||d | j                  ¬«      S ©NÚcall_function©ÚconstantÚtracer)ÚpytreeÚtree_mapr¸   Úunwrap_proxyÚcreate_proxyr   )
r±   r”   r   r€   ÚoutÚ
proxy_argsÚproxy_kwargsÚ	out_proxyr¯   r‡   s
           €€r   Úimpl_proxy_dispatch_modezDregister_run_and_save_rng_state_op.<locals>.impl_proxy_dispatch_modeÐ   s†   ø€ á! "Ð6 tÒ6¨vÑ6ˆÜ—_‘_ T§[¡[×%=Ñ%=À¸{ÀT¹{ÓKˆ
Ü—‘ t§{¡{×'?Ñ'?ÀÓHˆØ—K‘K×,Ñ,ØÐ3°ZÀó
ˆ	ô !  i¸$ÀtÇ{Á{ÔSÐSr   )r   Úpy_implr   ÚAutogradr   ÚCUDAÚCPUÚHPUÚXPUÚBackendSelectr
   r   )	r›   r²   rÁ   r¯   r£   r¡   r§   r©   r‡   s	      @@@@@@r   Ú"register_run_and_save_rng_state_oprÉ   ›   s@  ý€ ô9Ô0ô 9ñ 0Ó1Ðà8Ð×"Ñ"¤;×#7Ñ#7Ó8Ü Ð!7ÈÔMôð ×#Ñ#¤K×$4Ñ$4Ó5ñ?ó 6ð?ð ×#Ñ#¤K§O¡OÓ4ñ:ó 5ð:ð ×#Ñ#¤K§O¡OÓ4ñQó 5ðQð
 ×#Ñ#¤K§O¡OÓ4ñ>ó 5ð>ð ×#Ñ#¤K×$=Ñ$=Ó>ö)ó ?ð)ð ×#Ñ#¤NÓ3ó<ó 4ð<ð
 ×#Ñ#Ô$:Ó;ôTó <ðTð "Ð!r   c                  ó°  ‡‡‡‡‡	—  G d„ dt         «      }  | «       Š	 ‰	j                  t        j                  «      t	        ‰	d¬«      «       ‰	j                  t        j
                  «      d„ «       Š‰	j                  t        j                  «      d„ «       Š‰	j                  t        j                  «      d„ «       Š‰	j                  t        j                  «      d„ «       Š‰	j                  t        «      ˆ	fd	„«       }‰	j                  t        j                  «      ˆˆˆˆfd
„«       }‰	j                  t        «      d„ «       }‰	j                  ˆ	fd„«       }‰	S )Nc                   ó(   ‡ — e Zd Zˆ fd„Zˆ fd„Zˆ xZS )ú7register_run_with_rng_state_op.<locals>.RunWithRngStatec                 ó(   •— t         ‰| �  dd¬«       y )NÚrun_with_rng_stateTrˆ   rŠ   r�   s    €r   rŒ   z@register_run_with_rng_state_op.<locals>.RunWithRngState.__init__ß   ó   ø€ Ü‰GÑÐ1¸TÐÕBr   c                 ó,   •— t        ‰| �  ||g|¢­i |¤ŽS r‘   r’   )rŽ   Ú	rng_stater”   r   r€   r�   s        €r   r“   z@register_run_with_rng_state_op.<locals>.RunWithRngState.__call__â   s   ø€ ä‘7Ñ# I¨rÐC°DÒC¸FÑCÐCr   r•   rš   s   @r   ÚRunWithRngStaterÌ   Þ   s   ø„ ô	C÷	Dð 	Dr   rÒ   Trœ   c                 óê   — t         j                  j                  «       }t         j                  j                  | j	                  «       «        ||i |¤Ž}t         j                  j                  |«       |S r‘   )r   rB   rŸ   Úset_rng_statera   ©rÑ   r”   r   r€   Úcurrent_stater½   s         r   r¡   z1register_run_with_rng_state_op.<locals>.impl_cudaì   sR   € äŸ
™
×0Ñ0Ó2ˆÜ�
‰
× Ñ  §¡£Ô1Ù�$Ð!˜&Ñ!ˆÜ�
‰
× Ñ  Ô/Øˆ
r   c                 ó’   — t        j                  «       }t        j                  | «        ||i |¤Ž}t        j                  |«       |S r‘   )r   rŸ   rÔ   rÕ   s         r   r£   z0register_run_with_rng_state_op.<locals>.impl_cpuô   s@   € ä×+Ñ+Ó-ˆÜ×Ñ˜IÔ&Ù�$Ð!˜&Ñ!ˆÜ×Ñ˜MÔ*Øˆ
r   c                 ó  — t        t        d«      rft        j                  j                  «       }t        j                  j	                  | «        ||i |¤Ž}t        j                  j	                  |«       |S t        d«      ‚r¥   )r¦   r   ry   rŸ   rÔ   r   rÕ   s         r   r§   z0register_run_with_rng_state_op.<locals>.impl_hpuü   sb   € ä”5˜%Ô Ü!ŸI™I×3Ñ3Ó5ˆMÜ�I‰I×#Ñ# IÔ.Ù�dÐ%˜fÑ%ˆCÜ�I‰I×#Ñ# MÔ2ØˆJÜÐOÓPÐPr   c                 óÎ   — t         j                  j                  «       }t         j                  j                  | «        ||i |¤Ž}t         j                  j                  |«       |S r‘   )r   rv   rŸ   rÔ   rÕ   s         r   r©   z0register_run_with_rng_state_op.<locals>.impl_xpu  sL   € äŸ	™	×/Ñ/Ó1ˆÜ�	‰	×Ñ 	Ô*Ù�$Ð!˜&Ñ!ˆÜ�	‰	×Ñ Ô.Øˆ
r   c                 ó|  •— t        «       5   ‰	||g|¢­i |¤Ž}d d d «       t        j                  | j                  j                  ||g|¢­«      }t        j                  | j                  j                  |«      }| j                  j                  d‰	||«      }t        |d | j                  ¬«      S # 1 sw Y   Œ™xY wr´   ©r   r¹   rº   r¸   r»   r¼   r   )
r±   rÑ   r”   r   r€   r½   r¾   r¿   rÀ   rÎ   s
            €r   rÁ   z@register_run_with_rng_state_op.<locals>.impl_proxy_dispatch_mode  s®   ø€ ô )Ó*ñ 	EÙ$ Y°ÐD°TÒD¸VÑDˆC÷	Eä—_‘_ T§[¡[×%=Ñ%=À	È2Ð?UÐPTÑ?UÓVˆ
Ü—‘ t§{¡{×'?Ñ'?ÀÓHˆØ—K‘K×,Ñ,ØÐ/°¸\ó
ˆ	ô !  i¸$ÀtÇ{Á{ÔSÐS÷	Eð 	Eús   ŒB2Â2B;c                 ór   •— ‰‰‰	‰
dœ}t        ||«      }||vrt        d|› �«      ‚||   } || |g|¢­i |¤ŽS r«   r¬   )rÑ   r”   r   r€   r­   r   r®   r£   r¡   r§   r©   s          €€€€r   r¯   z;register_run_with_rng_state_op.<locals>.impl_backend_select  sb   ø€ ð ØØØñ	
ˆô ˜D &Ó)ˆØ˜Ñ!Ü Ð#=¸f¸XÐ!FÓGÐGØ˜ÑˆÙ�I˜rÐ3 DÒ3¨FÑ3Ð3r   c                 óB   — | 5   ||i |¤Žcd d d «       S # 1 sw Y   y xY wr‘   r   )r±   rÑ   r”   r   r€   s        r   r²   z=register_run_with_rng_state_op.<locals>.impl_fake_tensor_mode)  ó)   € ð ñ 	'Ù�tÐ&˜vÑ&÷	'÷ 	'ò 	'úó   ƒ•c                 óò   •— | j                  |«      }| j                  |«      }| j                  |«      }| j                  «       5   ‰	||g|¢­i |¤Ž}| j                  |«      cd d d «       S # 1 sw Y   y xY wr‘   ©Úunwrap_tensorsÚredispatch_to_nextÚwrap_tensors)
ÚctxrÑ   r”   r   r€   Úunwrapped_rng_stateÚunwrapped_argsÚunwrapped_kwargsr½   rÎ   s
            €r   Úimpl_functionalz7register_run_with_rng_state_op.<locals>.impl_functional0  s…   ø€ à!×0Ñ0°Ó;ÐØ×+Ñ+¨DÓ1ˆØ×-Ñ-¨fÓ5Ðà×#Ñ#Ó%ñ 	)Ù$Ø# RðØ*8òØ<LñˆCð ×#Ñ# CÓ(÷		)÷ 	)ò 	)ús   ÁA-Á-A6)r   rÂ   r   rÃ   r   rÄ   rÅ   rÆ   rÇ   r   rÈ   r
   Úpy_functionalize_impl)
rÒ   rÁ   r¯   r²   ré   r£   r¡   r§   r©   rÎ   s
        @@@@@r   Úregister_run_with_rng_state_oprë   Ý   s\  ü€ ôDÔ-ô Dñ )Ó*Ðà4Ð×Ñœ{×3Ñ3Ó4Ü Ð!3ÀDÔIôð ×Ñ¤× 0Ñ 0Ó1ñó 2ðð ×Ñ¤§¡Ó0ñó 1ðð ×Ñ¤§¡Ó0ñQó 1ðQð ×Ñ¤§¡Ó0ñó 1ðð ×ÑÔ 6Ó7ó
Tó 8ð
Tð ×Ñ¤× 9Ñ 9Ó:ö4ó ;ð4ð ×Ñ¤Ó/ñ'ó 0ð'ð ×-Ñ-ó	)ó .ð	)ð Ðr   c                  óä  ‡‡—  G d„ dt         «      }  | «       Š ‰j                  t        j                  «      t	        ‰d¬«      «       ‰j                  t        j
                  «      d dœd„«       Š‰j                  t        j                  «      d dœˆfd„
«       }‰j                  t        «      d dœd„«       }‰j                  t        «      d dœˆfd	„
«       }‰j                  d dœˆfd
„
«       }‰S )Nc                   ó.   ‡ — e Zd Zˆ fd„Zddœˆ fd„
Zˆ xZS )úJregister_graphsafe_run_with_rng_state_op.<locals>.GraphSafeRunWithRngStatec                 ó$   •— t         ‰| �  d«       y )NÚgraphsafe_run_with_rng_staterŠ   r�   s    €r   rŒ   zSregister_graphsafe_run_with_rng_state_op.<locals>.GraphSafeRunWithRngState.__init__E  s   ø€ Ü‰GÑÐ;Õ<r   N©rÑ   c                ó.   •— t        ‰| �  |g|¢­d|i|¤ŽS ©NrÑ   r’   )rŽ   r”   rÑ   r   r€   r�   s        €r   r“   zSregister_graphsafe_run_with_rng_state_op.<locals>.GraphSafeRunWithRngState.__call__H  s"   ø€ ä‘7Ñ# BÐM¨ÒM¸ÐMÀfÑMÐMr   r•   rš   s   @r   ÚGraphSafeRunWithRngStaterî   D  s   ø„ ô	=ð 15÷ 	Nò 	Nr   rô   Trœ   rñ   c                óà   — |j                   j                  }t        j                  j                  |   }|j                  «       }|j                  |«        | |i |¤Ž}|j                  |«       |S r‘   )r   Úindexr   rB   Údefault_generatorsÚgraphsafe_get_stateÚgraphsafe_set_state)r”   rÑ   r   r€   Ú
device_idxÚ	generatorrÖ   r½   s           r   r¡   z;register_graphsafe_run_with_rng_state_op.<locals>.impl_cudaR  sg   € ð ×%Ñ%×+Ñ+ˆ
Ü—J‘J×1Ñ1°*Ñ=ˆ	Ø!×5Ñ5Ó7ˆà×%Ñ% iÔ0Ù�$Ð!˜&Ñ!ˆØ×%Ñ% mÔ4Øˆ
r   c                ó^   •— t        ||«      }|dk7  rt        d|› �«      ‚ ‰| g|¢­d|i|¤ŽS )NrB   z6GraphSafe RNG operations only supported for CUDA, got rÑ   r¬   )r”   rÑ   r   r€   r   r¡   s        €r   r¯   zEregister_graphsafe_run_with_rng_state_op.<locals>.impl_backend_select^  sJ   ø€ ä˜D &Ó)ˆØ�VÒÜ ØHÈÈÐQóð ñ ˜ÐB˜dÒB¨iÐB¸6ÑBÐBr   c                óB   — | 5   ||i |¤Žcd d d «       S # 1 sw Y   y xY wr‘   r   )r±   r”   rÑ   r   r€   s        r   r²   zGregister_graphsafe_run_with_rng_state_op.<locals>.impl_fake_tensor_modeg  s'   € àñ 	'Ù�tÐ&˜vÑ&÷	'÷ 	'ò 	'úrß   c                ó„  •— t        «       5   ‰	|g|¢­d|i|¤Ž}d d d «       t        j                  | j                  j                  |g|¢­«      }t        j                  | j                  j                  d|i|¥«      }| j                  j                  d‰	||«      }t        |d | j                  ¬«      S # 1 sw Y   ŒœxY w)NrÑ   rµ   r¶   rÛ   )
r±   r”   rÑ   r   r€   r½   r¾   r¿   rÀ   rð   s
            €r   rÁ   zJregister_graphsafe_run_with_rng_state_op.<locals>.impl_proxy_dispatch_model  sº   ø€ ä(Ó*ñ 	YÙ.¨rÐX°DÒXÀIÐXÐQWÑXˆC÷	Yä—_‘_ T§[¡[×%=Ñ%=À¸{ÀT¹{ÓKˆ
Ü—‘Ø�K‰K×$Ñ$ {°IÐ&HÀÐ&Hó
ˆð —K‘K×,Ñ,ØÐ9¸:À|ó
ˆ	ô !  i¸$ÀtÇ{Á{ÔSÐS÷	Yð 	Yús   ŒB6Â6B?c                óü   •— |�| j                  |«      nd }| j                  |«      }| j                  |«      }| j                  «       5   ‰	|g|¢­d|i|¤Ž}| j                  |«      cd d d «       S # 1 sw Y   y xY wró   rá   )
rå   r”   rÑ   r   r€   ræ   rç   rè   r½   rð   s
            €r   ré   zAregister_graphsafe_run_with_rng_state_op.<locals>.impl_functionaly  s˜   ø€ ð .7Ð-BˆC×Ñ˜yÔ)Èð 	ð ×+Ñ+¨DÓ1ˆØ×-Ñ-¨fÓ5Ðà×#Ñ#Ó%ñ 	)Ù.ØðØ#òØ/BðØFVñˆCð ×#Ñ# CÓ(÷		)÷ 	)ò 	)ús   Á	A2Á2A;©
r   rÂ   r   rÃ   r   rÄ   rÈ   r
   r   rê   )rô   r¯   r²   rÁ   ré   rð   r¡   s        @@r   Ú(register_graphsafe_run_with_rng_state_opr  C  s  ù€ ôNÔ#6ô Nñ $<Ó#=Ð à>Ð ×(Ñ(¬×)=Ñ)=Ó>Ü Ð!=ÈdÔSôð "×)Ñ)¬+×*:Ñ*:Ó;Ø'+ó 	ó <ð	ð "×)Ñ)¬+×*CÑ*CÓDØ15õ Có EðCð "×)Ñ)¬.Ó9Ø9=ó 'ó :ð'ð "×)Ñ)Ô*@ÓAØ<@õ 
Tó Bð
Tð "×7Ñ7Ø26õ )ó 8ð)ð (Ð'r   c                  óÆ  ‡‡—  G d„ dt         «      }  | «       Š ‰j                  t        j                  «      t	        ‰d¬«      «       ‰j                  t        j
                  «      d„ «       Š‰j                  t        j                  «      ˆfd„«       }‰j                  t        «      d„ «       }‰j                  t        «      ˆfd„«       }‰j                  ˆfd	„«       }‰S )
u®  
    Register a higher-order operator for DTensor distributed random operations.
    Takes pre-computed integer offsets (start_offset_incr, end_offset_incr), an op,
    and args. Internally adjusts the RNG state using the offsets before/after
    running the op.

    The offsets are computed at trace time from the DTensorSpec, so the compiled graph
    contains only plain integers â€” no DTensorSpec or DeviceMesh objects.
    c                   ó(   ‡ — e Zd Zˆ fd„Zˆ fd„Zˆ xZS )ú4register_run_dtensor_rng_op.<locals>.RunDTensorRngOpc                 ó(   •— t         ‰| �  dd¬«       y )NÚrun_dtensor_rng_opTrˆ   rŠ   r�   s    €r   rŒ   z=register_run_dtensor_rng_op.<locals>.RunDTensorRngOp.__init__™  rÏ   r   c                 ó.   •— t        ‰| �  |||g|¢­i |¤ŽS r‘   r’   )rŽ   Ústart_offset_incrÚend_offset_incrr”   r   r€   r�   s         €r   r“   z=register_run_dtensor_rng_op.<locals>.RunDTensorRngOp.__call__œ  s,   ø€ ä‘7Ñ#Ø! ?°BðØ9=òØAGñð r   r•   rš   s   @r   ÚRunDTensorRngOpr  ˜  s   ø„ ô	C÷	ð 	r   r
  Trœ   c                 ó¢  — ddl m}  |t        j                  j	                  «       «      }|j
                  j                  «       }|| z   |_        |rt        |d   d«      r|d   j                  n3t        j                  dt        j                  j                  «       › �«      }t        j                  j                  |gd¬«      5  t        j                  j                  |j                  «       	  ||i |¤Ž}	||z   |_        	 d d d «       t        j                  j                  |j                  «       	S # ||z   |_        w xY w# 1 sw Y   ŒCxY w)Nr   )Ú_PhiloxStater   zcuda:rB   )rf   Údevice_type)Ú torch.distributed.tensor._randomr  r   rB   rŸ   rW   Úcloner¦   r   rD   rb   rc   rÔ   Ústate)
r  r	  r”   r   r€   r  r  Ú
old_offsetr   r½   s
             r   r¡   z.register_run_dtensor_rng_op.<locals>.impl_cuda¨  s  € åAáœUŸZ™Z×5Ñ5Ó7Ó8ˆØ—\‘\×'Ñ'Ó)ˆ
Ø!Ð$5Ñ5ˆŒñ œ  Q¡¨Ô2ð �‰G�NŠNä—‘ ¤e§j¡j×&?Ñ&?Ó&AÐ%BÐCÓDð 	ô
 �\‰\×"Ñ"¨F¨8ÀÐ"ÓHñ 	<Ü�J‰J×$Ñ$ U§[¡[Ô1ð<Ù˜$Ð) &Ñ)�à)¨OÑ;�•÷	<ô 	�
‰
× Ñ  §¡Ô-Øˆ
øð  *¨OÑ;�•ú÷	<ð 	<ús$   Ã*EÃ0D6Ã8
EÄ6EÅEÅEc                 ó`   •— t        ||«      }|dk7  rt        d|› d�«      ‚ ‰| ||g|¢­i |¤ŽS )NrB   z2run_dtensor_rng_op only supports CUDA device, got zF. This operator is designed for distributed random operations on CUDA.)r‚   r   )r  r	  r”   r   r€   r   r¡   s         €r   r¯   z8register_run_dtensor_rng_op.<locals>.impl_backend_select¿  sT   ø€ ä˜D &Ó)ˆØ�VÒÜØDÀVÀHð MWð Xóð ñ Ð*¨O¸RÐQÀ$ÒQÈ&ÑQÐQr   c                 óB   — | 5   ||i |¤Žcd d d «       S # 1 sw Y   y xY wr‘   r   )r±   r  r	  r”   r   r€   s         r   r²   z:register_run_dtensor_rng_op.<locals>.impl_fake_tensor_modeÉ  rÞ   rß   c                 ó€  •— t        «       5   ‰
|||g|¢­i |¤Ž}d d d «       t        j                  | j                  j                  |||g|¢­«      }t        j                  | j                  j                  |«      }| j                  j                  d‰
||«      }	t        |	d | j                  ¬«      S # 1 sw Y   ŒšxY wr´   rÛ   )r±   r  r	  r”   r   r€   r½   r¾   r¿   rÀ   r  s             €r   rÁ   z=register_run_dtensor_rng_op.<locals>.impl_proxy_dispatch_modeÐ  s¿   ø€ ô )Ó*ñ 	Ù$Ø! ?°BðØ9=òØAGñˆC÷	ô —_‘_Ø�K‰K×$Ñ$Ð'8¸/È2Ð&UÐPTÑ&Uó
ˆ
ô —‘ t§{¡{×'?Ñ'?ÀÓHˆØ—K‘K×,Ñ,ØÐ/°¸\ó
ˆ	ô !  i¸$ÀtÇ{Á{ÔSÐS÷	ð 	ús   ŒB4Â4B=c                 óÒ   •— | j                  |«      }| j                  |«      }| j                  «       5   ‰	|||g|¢­i |¤Ž}| j                  |«      cd d d «       S # 1 sw Y   y xY wr‘   rá   )
rå   r  r	  r”   r   r€   rç   rè   r½   r  s
            €r   ré   z4register_run_dtensor_rng_op.<locals>.impl_functionalá  s|   ø€ à×+Ñ+¨DÓ1ˆØ×-Ñ-¨fÓ5Ðà×#Ñ#Ó%ñ 	)Ù$Ø!ØØðð  ò	ð
 #ñˆCð ×#Ñ# CÓ(÷	)÷ 	)ò 	)ús   ´AÁA&r   )r
  r¯   r²   rÁ   ré   r¡   r  s        @@r   Úregister_run_dtensor_rng_opr  �  sñ   ù€ ôÔ-ô ñ )Ó*Ðà4Ð×Ñœ{×3Ñ3Ó4Ü Ð!3ÀDÔIôð ×Ñ¤× 0Ñ 0Ó1ñó 2ðð, ×Ñ¤× 9Ñ 9Ó:óRó ;ðRð ×Ñ¤Ó/ñ'ó 0ð'ð ×ÑÔ 6Ó7óTó 8ðTð  ×-Ñ-ó)ó .ð)ð Ðr   c                  ó   — t        «        y r‘   )rm   r   r   r   Úregister_rng_primsr  ö  s   € ÜÕr   r‘   )*Útypingr   r   Útorch.utils._pytreeÚutilsÚ_pytreer¹   r   Útorch._Cr   Útorch._higher_order_ops.utilsr   Ú
torch._opsr   Útorch._prims_commonr   r	   Útorch._subclasses.fake_tensorr
   Ú"torch.fx.experimental.proxy_tensorr   r   r   Útorch.typesr   r   r   r4   rh   r<   rS   rm   r‚   rÉ   rë   r‡   rÎ   r  rð   r  r  r  r   r   r   ú<module>r$     s¸   ðå ã ß $Ð $Ý Ý  Ý BÝ *ß OÝ 8÷ñ ÷
 (òó%ð.AØ�:‰:óAðWØ�:‰:óWò,6òrò&?"òD_ñD <Ó=Ð Ù3Ó5Ð òD(ñN  HÓIÐ òcñL 1Ó2Ð ór   