Ë
    ñukiÿ  ã                  ó¤  — U d Z ddlmZ ddlmZ  G d„ de«      Z G d„ de«      Z G d„ d	e«      Z G d
„ de«      Zg a	de
d<   g ade
d<   g ade
d<   g Zde
d<   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d0d„Zd1d„Zd „ Z	 	 	 	 d,d!„Z	 	 	 	 d+d"„Z	 	 	 	 d*d#„Z	 	 	 	 d-d$„Zy%)2a  Utilities for instrumenting code.

Code points can be marked as a named event. Every time an event is reached
during program execution, the registered listeners will be invoked.

A typical listener callback is to send an event to a metrics collector for
aggregation/exporting.
é    )Úannotations)ÚProtocolc                  ó   — e Zd Zdd„Zy)ÚEventListenerWithMetadatac                 ó   — y ©N© )ÚselfÚeventÚkwargss      úN/home/cdr/jupyterlab/.venv/lib/python3.12/site-packages/jax/_src/monitoring.pyÚ__call__z"EventListenerWithMetadata.__call__   s   € Øó    N©r   Ústrr   ú	str | intÚreturnÚNone©Ú__name__Ú
__module__Ú__qualname__r   r	   r   r   r   r      s   „ ôr   r   c                  ó   — e Zd Z	 	 	 	 dd„Zy)Ú!EventDurationListenerWithMetadatac                 ó   — y r   r	   )r
   r   Úduration_secsr   s       r   r   z*EventDurationListenerWithMetadata.__call__%   s   € àr   N)r   r   r   Úfloatr   r   r   r   r   r	   r   r   r   r   #   s   „ ðØ"ðØ'+ôr   r   c                  ó(   — e Zd Z	 	 	 	 	 	 	 	 	 	 dd„Zy)Ú!EventTimeSpanListenerWithMetadatac                 ó   — y r   r	   )r
   r   Ú
start_timeÚend_timer   s        r   r   z*EventTimeSpanListenerWithMetadata.__call__,   ó   € ð r   N©
r   r   r!   r   r"   r   r   r   r   r   r   r	   r   r   r   r   *   s/   „ ðØðØ$)ðØ5:ðØFOðàôr   r   c                  ó$   — e Zd Z	 	 	 	 	 	 	 	 dd„Zy)ÚScalarListenerWithMetadatac                 ó   — y r   r	   )r
   r   Úvaluer   s       r   r   z#ScalarListenerWithMetadata.__call__3   r#   r   N©r   r   r(   zfloat | intr   r   r   r   r   r	   r   r   r&   r&   1   s'   „ ðØðØ*ðØ6?ðàôr   r&   úlist[EventListenerWithMetadata]Ú_event_listenersú'list[EventDurationListenerWithMetadata]Ú_event_duration_secs_listenersú'list[EventTimeSpanListenerWithMetadata]Ú_event_time_span_listenersú list[ScalarListenerWithMetadata]Ú_scalar_listenersc                ó,   — t         D ]  } || fi |¤Ž Œ y)z¬Record an event.

  If **kwargs are specified, all of the named arguments have to be passed in the
  same order across all invocations of this method for the same event.
  N)r+   )r   r   Úcallbacks      r   Úrecord_eventr4   ?   s    € ô #ò €hÙˆUÑfÓñr   c                ó.   — t         D ]  } || |fi |¤Ž Œ y)zÈRecord an event duration in seconds (float).

  If **kwargs are specified, all of the named arguments have to be passed in the
  same order across all invocations of this method for the same event.
  N)r-   )r   Údurationr   r3   s       r   Úrecord_event_duration_secsr7   I   s"   € ô 1ò (€hÙˆUHÑ' Ó'ñ(r   c                ó0   — t         D ]  } || ||fi |¤Ž Œ y)z6Record an event start and end time in seconds (float).N)r/   )r   r!   r"   r   r3   s        r   Úrecord_event_time_spanr9   T   s$   € ô -ò 4€hÙˆUJ Ñ3¨FÓ3ñ4r   c                ó.   — t         D ]  } || |fi |¤Ž Œ y)zRecord a scalar summary value.N©r1   )r   r(   r   r3   s       r   Úrecord_scalarr<   \   s"   € ô $ò %€hÙˆUEÑ$˜VÓ$ñ%r   c                ó.   — t         j                  | «       y)z8Register a callback to be invoked during record_event().N)r+   Úappend©r3   s    r   Úregister_event_listenerr@   d   s   € ô ×Ñ˜(Õ#r   c                ó.   — t         j                  | «       y)zBRegister a callback to be invoked during record_event_time_span().N)r/   r>   r?   s    r   Ú!register_event_time_span_listenerrB   k   s   € ô ×#Ñ# HÕ-r   c                ó.   — t         j                  | «       y)zFRegister a callback to be invoked during record_event_duration_secs().N)r-   r>   r?   s    r   Ú%register_event_duration_secs_listenerrD   r   s   € ô !×'Ñ'¨Õ1r   c                ó.   — t         j                  | «       y)z9Register a callback to be invoked during record_scalar().N)r1   r>   r?   s    r   Úregister_scalar_listenerrF   x   s   € ô ×Ñ˜8Õ$r   c                 ó    — t        t        «      S )zGet event duration listeners.)Úlistr-   r	   r   r   Úget_event_duration_listenersrI      s   € ä	Ô,Ó	-Ð-r   c                 ó    — t        t        «      S )zGet event time span listeners.)rH   r/   r	   r   r   Úget_event_time_span_listenersrK   „   s   € ä	Ô(Ó	)Ð)r   c                 ó    — t        t        «      S )zGet event listeners.)rH   r+   r	   r   r   Úget_event_listenersrM   ‰   s   € ä	ÔÓ	Ðr   c                 ó    — t        t        «      S )zGet scalar event listeners.)rH   r1   r	   r   r   Úget_scalar_listenersrO   Ž   s   € ä	ÔÓ	 Ð r   c                 ó   — g a g ag ag } y)zClear event listeners.N)r+   r-   r/   r;   s    r   Úclear_event_listenersrQ   “   s   € ð ÐØ#%Ð Ø!ÐØÑr   c                óB   — | t         v sJ ‚t         j                  | «       y)z2Unregister an event duration listener by callback.N)r-   Úremover?   s    r   Ú"unregister_event_duration_listenerrT   œ   s!   € ð 
Ô3Ñ	3Ð3Ð	3Ü ×'Ñ'¨Õ1r   c                óB   — | t         v sJ ‚t         j                  | «       y)z3Unregister an event time span listener by callback.N)r/   rS   r?   s    r   Ú#unregister_event_time_span_listenerrV   ¤   s!   € ð 
Ô/Ñ	/Ð/Ð	/Ü×#Ñ# HÕ-r   c                óB   — | t         v sJ ‚t         j                  | «       y)z)Unregister an event listener by callback.N)r+   rS   r?   s    r   Úunregister_event_listenerrX   ¬   s!   € ð 
Ô%Ñ	%Ð%Ð	%Ü×Ñ˜(Õ#r   c                óB   — | t         v sJ ‚t         j                  | «       y)z/Unregister a scalar event listener by callback.N)r1   rS   r?   s    r   Úunregister_scalar_listenerrZ   ´   s!   € ð 
Ô&Ñ	&Ð&Ð	&Ü×Ñ˜8Õ$r   Nr   )r   r   r6   r   r   r   r   r   r$   r)   )r3   r   r   r   )r3   r   r   r   )r3   r   r   r   )r3   r&   r   r   )r   r,   )r   r.   )r   r*   )r   r0   )Ú__doc__Ú
__future__r   Útypingr   r   r   r   r&   r+   Ú__annotations__r-   r/   r1   r4   r7   r9   r<   r@   rB   rD   rF   rI   rK   rM   rO   rQ   rT   rV   rX   rZ   r	   r   r   ú<module>r_      s¡  ðòõ #å ô ô ô¨ô ô¨ô ô ô ð 57Ð Ð1Ó 6ØJLÐ Ð GÓ LØFHÐ ÐCÓ HØ68Ð Ð3Ó 8óð(Ø)2ð(Ø7;ó(ð4Øð4Ø!ð4Ø-2ð4Ø>Gð4à	ó4ð%Øð%Ø"ð%Ø.7ð%à	ó%ð$Ø'ð$à	ó$ð.Ø/ð.à	ó.ð2Ø0ð2Ø59ó2ð%Ø)ð%à	ó%ó.ó
*ó
 ó
!ò
ð2Ø/ð2à	ó2ð.Ø/ð.à	ó.ð$Ø'ð$à	ó$ð%Ø(ð%à	ô%r   