
      i	                         d dl mZ d dlmZmZ d dlZd dlmZ dedeeej                  f   defdZ	 G d	 d
      Z
 G d d      Zy)    )Mapping)AnyUnionN)Tensortensordst_typereturnc                 R    t        j                  |       r| j                  |      S | S N)torchis_floating_pointto)r   r   s     }/Volumes/fast/ai/experiments/voice-extract-mac/.venv/lib/python3.12/site-packages/lightning/fabric/plugins/precision/utils.py_convert_fp_tensorr      s#    "'"9"9&"A699XMvM    c                   T    e Zd ZdZdej
                  ddfdZddZdeded	eddfd
Z	y)_DtypeContextManagerz}A context manager to change the default tensor type when tensors get created.

    See: :func:`torch.set_default_dtype`

    dtyper	   Nc                 D    t        j                         | _        || _        y r   )r   get_default_dtype_previous_dtype
_new_dtype)selfr   s     r   __init__z_DtypeContextManager.__init__!   s    ,1,C,C,Er   c                 B    t        j                  | j                         y r   )r   set_default_dtyper   )r   s    r   	__enter__z_DtypeContextManager.__enter__%   s    0r   exc_type	exc_value	tracebackc                 B    t        j                  | j                         y r   )r   r   r   )r   r   r   r    s       r   __exit__z_DtypeContextManager.__exit__(   s     4 45r   r	   N)
__name__
__module____qualname____doc__r   r   r   r   r   r"    r   r   r   r      sD     ekk  d  16 6 6 6 6r   r   c                   J    e Zd ZdZdeeef   ddfdZddZde	de	d	e	ddfd
Z
y)_ClassReplacementContextManagerz)A context manager to monkeypatch classes.mappingr	   Nc                     || _         i | _        i | _        |D ]M  }|j                  dd      \  }}t	        ||g      }|| j                  |<   t        ||      | j                  |<   O y )N.   )fromlist)_mapping
_originals_modulesrsplit
__import__getattr)r   r+   class_stringmodule_name
class_namemodules         r   r   z(_ClassReplacementContextManager.__init__/   si    #L&2&9&9#q&A#KzlCF*0DMM,',3FJ,GDOOL)	 $r   c                     | j                   j                         D ]4  \  }}|j                  dd      \  }}t        | j                  |   ||       6 y Nr-   r.   )r0   itemsr3   setattrr2   )r   r6   replacement_r8   s        r   r   z)_ClassReplacementContextManager.__enter__9   sI    )-)<)<)>%L+(//Q7MAzDMM,/[I *?r   r   r   r    c                     | j                   j                         D ]A  \  }}|j                  dd      \  }}t        | j                  |   || j
                  |          C y r;   )r0   r<   r3   r=   r2   r1   )r   r   r   r    r6   r>   r?   r8   s           r   r"   z(_ClassReplacementContextManager.__exit__>   sR    )-)<)<)>%L+(//Q7MAzDMM,/T__\=Z[ *?r   r#   )r$   r%   r&   r'   r   strtyper   r   r   r"   r(   r   r   r*   r*   ,   sP    3HT	 2 Ht HJ
\ \ \ \ \r   r*   )collections.abcr   typingr   r   r   r   rA   r   r   r   r*   r(   r   r   <module>rE      sU    $   Nv NsEKK7G1H NV N
6 6$\ \r   