
    iX"                     D   d Z ddlZddlZddlZddlZddlZddlZddlZddlZddl	m
Z
 ddlmZmZ ddlZddlZddlmZ ddlmZ ddlmZmZmZmZ dd	lmZmZmZ  ej<                  e      Z dd
ee!e"e!df   f   ddfdZ# G d de      Z G d de      Z$ G d de      Z%y)a  Testing utilities for Dynamo, providing a specialized TestCase class and test running functionality.

This module extends PyTorch's testing framework with Dynamo-specific testing capabilities.
It includes:
- A custom TestCase class that handles Dynamo-specific setup/teardown
- Test running utilities with dependency checking
- Automatic reset of Dynamo state between tests
- Proper handling of gradient mode state
    N)Callable)AnyUnion)	polyfills)	trace_log)
IS_WINDOWSTEST_WITH_CROSSREFTEST_WITH_TORCHDYNAMOTestCase   )configresetutilsneeds.returnc                    ddl m} t        st        ry t        j
                  j                         s*t        r$t        j                  j                  dd      dk(  ry t        | t              r| f} | D ]>  }|dk(  r!t        j                  j                         r' y 	 t        j                  |       @  |        y # t         $ r Y  y w xY w)Nr   )	run_testsTORCHINDUCTOR_WINDOWS_TESTS0cuda)$torch.testing._internal.common_utilsr   r
   r	   torchxpuis_availabler   osenvironget
isinstancestrr   	importlibimport_moduleImportError)r   r   needs      h/var/www/vps2.regionflexible.com/Desarrollo/venv/lib/python3.12/site-packages/torch/_dynamo/test_case.pyr   r   '   s    > 2 II""$JJNN8#>#E% 6>::**,''- K  s   B::	CCc            
            e Zd ZU ej                  ed<   ed fd       Zed fd       Zd fdZ	d fdZ
ded	ed
ededdf
 fdZ xZS )r   _exit_stackr   Nc                 V    | j                   j                          t        |           y N)r&   closesupertearDownClasscls	__class__s    r$   r+   zTestCase.tearDownClassE   s        c                     t         |           t        j                         | _        | j                  j                  t        j                  ddd             y )NTF)raise_on_ctx_manager_usagesuppress_errorslog_compilation_metrics)r*   
setUpClass
contextlib	ExitStackr&   enter_contextr   patchr,   s    r$   r4   zTestCase.setUpClassJ   sE    $..0%%LL+/ %(-	
r/   c                    t        j                         | _        t        |           t                t        j                  j                          t        j                         | _        t        j                  | j                         y r(   )r   is_grad_enabled_prior_is_grad_enabledr*   setUpr   r   countersclearloggingNullHandlerhandlerr   
addHandlerselfr.   s    r$   r<   zTestCase.setUpV   sT    &+&;&;&=#**,T\\*r/   c                 *   t        j                  | j                         t        j                  j                         D ]  \  }}t        ||j                                ! t                t        j                  j                          t        j                  j                  j                          t        | =          | j                   t        j"                         ur5t$        j'                  d       t        j(                  | j                          y y )NzRunning test changed grad mode)r   removeHandlerrA   r   r=   itemsprintmost_commonr   r>   r   _C	_autograd_saved_tensors_hooks_enabler*   tearDownr;   r:   logwarningset_grad_enabled)rD   kvr.   s      r$   rM   zTestCase.tearDown^   s    -NN((* 	&DAq!Q]]_%	&668&&e.C.C.EEKK89""4#>#>? Fr/   xyargskwargsc                     t         j                  rt        |t        j                        st        |t        j                        ry t        |   ||g|i |S r(   )r   debug_disable_compile_counterr   r   CompileCounterIntr*   assertEqual)rD   rS   rT   rU   rV   r.   s        r$   rZ   zTestCase.assertEqualj   sK    001e556!U445w"1a9$9&99r/   r   N)__name__
__module____qualname__r5   r6   __annotations__classmethodr+   r4   r<   rM   r   rZ   __classcell__r.   s   @r$   r   r   B   sm    %%%    	
 	
+
@:S :S : : : : :r/   r   c                   ,     e Zd Zd fdZd fdZ xZS )TestCaseWithNestedGraphBreaksc                     t         |           t        j                  j                  j
                  | _        dt        j                  j                  _        y )NT)r*   r<   r   _dynamor   nested_graph_breaksprev_nested_graph_breaksrC   s    r$   r<   z#TestCaseWithNestedGraphBreaks.setUpz   s6    (-(<(<(P(P%370r/   c                 t    t         |           | j                  t        j                  j
                  _        y r(   )r*   rM   rh   r   rf   r   rg   rC   s    r$   rM   z&TestCaseWithNestedGraphBreaks.tearDown   s&    373P3P0r/   r[   )r\   r]   r^   r<   rM   ra   rb   s   @r$   rd   rd   y   s    8Q Qr/   rd   c            
       h    e Zd ZU dZej
                  ed<   dZej                  j                  Z
ej                  j                  Zej                  j                  Zej                  j                  Zej                  j                  Zej                  j                  Zej                  j                   Zej                  j"                  Zej                  j$                  Zej                  j&                  Zej                  j(                  Zej                  j*                  Zej                  j,                  Zej                  j.                  Zej                  j0                  Zej                  j2                  Zej                  j4                  Zej                  j6                  Zej                  j8                  Zej                  j:                  Zej                  j<                  Zej@                  Z!ejD                  Z#ej                  jH                  Z$ej                  jJ                  Z%ej                  jL                  Z&ejN                  Z(ej                  jR                  Z)ej                  jT                  Z*ej                  jV                  Z+ej                  jX                  Z,ej                  jZ                  Z-ej                  j\                  Z.ej                  j^                  Z/de0de1f   de2e3e0de1f   f   de4de0de1f   fd	Z5de3f fd
Z6e7d fd       Z8e7d fd       Z9de3de0ddfdZ: xZ;S )CPythonTestCasez
    Test class for CPython tests located in "test/dynamo/CPython/Py_version/*".

    This class enables specific features that are disabled by default, such as
    tracing through unittest methods.
    _stackTfn.backendnopythonr   c                     t        | | j                        }t        j                  j	                  ||      |      }t        | | j                  |       |S )N)error_on_graph_break)getattr_testMethodNamer   rf   optimizesetattr)rD   rm   rn   ro   methods        r$   
compile_fnzCPythonTestCase.compile_fn   sL     t334''h'OPVWd**F3	r/   c                    t         |          }| j                  }t        j                  |      j                  t        j                        d   j                  d      d   }t        j                  dt        j                  |            }|r>|j                         j                  t        j                        j                  dd      }n|S d| d| d| S )	N.r   z
/([\d_]+)/_ CPython-)r*   _dynamo_test_keyr.   inspectgetfilesplitr   sepresearchgroupstripreplace)rD   suffixtest_cls	test_filepy_verr.   s        r$   r   z CPythonTestCase._dynamo_test_key   s    )+>>OOH-33BFF;B?EEcJ1M	='//(*CD\\^))"&&199#rBFM)AfX66r/   Nc                 V    | j                   j                          t        |           y r(   )rl   r)   r*   r+   r,   s    r$   r+   zCPythonTestCase.tearDownClass   s    

r/   c                 t   t         j                  j                  dd      t         j                  j                  z   }t	        j
                  |      dz   }t        j                  |       }t	        j                  ||      }|rt        t        t        |j                         j                  |      j                  d                  }t        j                   d d }||k7  rdj                  t        t"        |            }dj                  t        t"        |            }t%        j&                  d| d|       t%        j&                  d	t        j                  |              t(        	| U          t-        j.                         | _        | j0                  j3                  t5        j6                  d
             y )Ndynamocpythonz\d_\d{2}r{      rz   zTest requires Python z but got Python z>Test requires a specific Python version but not found in path T)enable_trace_unittest)r   pathjoinr   r   escaper   r   r   tuplemapintr   removeprefixr   sysversion_infor   unittestSkipTestr*   r4   r5   r6   rl   r7   r   r8   )
r-   prefixregexsearch_pathmtest_py_verr   expectedgotr.   s
            r$   r4   zCPythonTestCase.setUpClass   sZ    h	2RWW[[@		&!K/ooc*IIe[)C)?)?)G)M)Mc)R STK%%bq)F$88C[$9:hhs3/0''+H:5EcUK  ##PQXQ`Q`adQePfg  	))+


  LL&*	
r/   method_namepolicyc                      y r(    )rD   r   r   s      r$   wrap_with_policyz CPythonTestCase.wrap_with_policy   s    r/   r[   )<r\   r]   r^   __doc__r5   r6   r_   dynamo_strict_nopythonr   r   rZ   assertNotEqual
assertTrueassertFalseassertIsassertIsNotassertIsNoneassertIsNotNoneassertInassertNotInassertIsInstanceassertNotIsInstanceassertAlmostEqualassertNotAlmostEqualassertGreaterassertGreaterEqual
assertLessassertLessEqualassertRegexassertNotRegexassertCountEqualr   assert_multi_line_equalassertMultiLineEqualassert_sequence_equalassertSequenceEqualassertListEqualassertTupleEqualassertSetEqualassert_dict_equalassertDictEqualassertRaisesassertRaisesRegexassertWarnsassertWarnsRegex
assertLogsfailfailureExceptionr   r   r   r   boolrw   r   r`   r+   r4   r   ra   rb   s   @r$   rk   rk      s       ! ##//K&&55N""--J##//K  ))H##//K$$11L''77O  ))H##//K((99"++?? ));;#,,AA%%33M!**==""--J''77O##//K&&55N((99$<<#99''77O((99&&55N11O$$11L ));;##//K((99""--J!!D((99S#X sHS#X../ 	
 
#s(		7# 	7     
 
:C  d r/   rk   )r   )&r   r5   r    r   r?   r   r   r   r   collections.abcr   typingr   r   r   torch.testingtorch._dynamor   torch._logging._internalr   r   r   r	   r
   r   TorchTestCaser|   r   r   r   	getLoggerr\   rN   r   r   r   rd   rk   r   r/   r$   <module>r      s        	 	 
  $    # .  # " g!U3c3h/0 $ 6/:} /:n
QH 
Qph pr/   