
    BVhP                         d 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 ddlm	Z	 ddl
mZ ddlmZ  G d d	ej                        Zy)
zTesting utilities.    N)def_function)op_callbacks)ops)	variables)testc                        e Zd ZdZ fdZ	 ddZd Zd Z fdZ fdZ	d Z
d	 Zd
 Zd Z fdZ fdZ fdZ fdZd Z xZS )AutoGraphTestCasea"  Tests specialized for AutoGraph, which run as tf.functions.

  These tests use a staged programming-like approach: most of the test code runs
  as-is inside a tf.function, but the assertions are lifted outside the
  function, and run with the corresponding function values instead.

  For example, the test:

      def test_foo(self):
        baz = bar();
        self.assertEqual(baz, value)

  is equivalent to writing:

      def test_foo(self):
        @tf.function
        def test_fn():
          baz = bar();
          return baz, value

        baz_actual, value_actual = test_fn()
        self.assertEqual(baz_actual, value_actual)

  Only assertions that require evaluation outside the function are lifted
  outside the function scope. The rest execute inline, at function creation
  time.
  c           	      0   t         |   |       }| j                  D ]v  }|j                  t        j
                  j                        s-t        ||      }t        |      sE|j                  |      }t        ||t        j                  ||             x |S N)super__new____dict__
startswithunittest
TestLoadertestMethodPrefixgetattrcallable_run_as_tf_functionsetattrtypes
MethodType)clsargsobjnamemwrapper	__class__s         Y/home/dcms/DCMS/lib/python3.12/site-packages/tensorflow/python/autograph/utils/testing.pyr   zAutoGraphTestCase.__new__:   s    
'/#
C ;__X00AAB
#t
a	!))!,T5++GS9:; J    c                 :    | j                   j                  |       y r   )	trace_logappend)selfop_typeinputsattrsoutputsop_namegraphs          r    _op_callbackzAutoGraphTestCase._op_callbackG   s    NN'"r!   c                     fd}|S )Nc                     t        j                  d       fd       }	  |       } j                  D ]"  } ||j                         j                         $  j                  |      }t         j                  |      D ]  \  \  }}} ||   y #   j                  -  j                  j                  t        j                           Y y  xY w)NF)	autographc                      g _         d _        g _        g _                 j                   D  cg c]  \  } }|	 }} }|S c c}} w r   )
assertions	raises_cmgraph_assertionsr#   )_r   targetsfnr%   s      r    
fn_wrapperzJAutoGraphTestCase._run_as_tf_function.<locals>.wrapper.<locals>.fn_wrapperN   sJ     "
'+7GAt477 8s   A)r   functionr3   get_concrete_functionr+   evaluater2   __exit__sysexc_infozipr1   )r%   r7   tensors	assertionactualsr4   valuesr6   s   `      r    r   z6AutoGraphTestCase._run_as_tf_function.<locals>.wrapperM   s    u- .,.. 	>I
J446<<
=	> --( %($A 
 .9a&6>>%
!$..
!
!3<<>
2

s   A	B 9CC )r%   r6   r   s    ` r    r   z%AutoGraphTestCase._run_as_tf_functionK   s    < Nr!   c                 &   t        j                         5  || j                  vrLt        j                  ||      | j                  |<   | j	                  | j                  |   j
                         d d d        | j                  |   S # 1 sw Y   xY w)N)dtype)r   
init_scoper   Variabler:   initializer)r%   r   valuerE   s       r    variablezAutoGraphTestCase.variablem   st    		 8	T^^	#(11%uEtdnnT*6678 >>$	8 8s   ABBc                     t         |           i | _        g | _        d | _        t        j                  | j                         y r   )r   setUpr   r#   r2   r   add_op_callbackr,   r%   r   s    r    rL   zAutoGraphTestCase.setUpt   s5    	GMODNDNDN  !2!23r!   c                 |    t        j                  | j                         d | _        d | _        t
        |           y r   )r   remove_op_callbackr,   r#   r   r   tearDownrN   s    r    rQ   zAutoGraphTestCase.tearDown{   s0    ##D$5$56DNDN	Gr!   c                 N      fd} j                   j                  |       y )Nc                    g }| j                         j                  D ]4  }t        j                  |j                        s$|j                  |       6 | j                         j                  j                  D ]E  }|j                  D ]4  }t        j                  |j                        s$|j                  |       6 G j                  |       y r   )
as_graph_defnoderematchr   r$   libraryr8   node_def	assertLen)r+   matchesrU   r6   rY   nop_regexr%   s        r    r@   z8AutoGraphTestCase.assertGraphContains.<locals>.assertion   s    g$$&++ $88Hdii(
..
 ""$,,55 %" 	%HXXh.NN8$	%% nnWa r!   )r3   r$   )r%   r]   r\   r@   s   ``` r    assertGraphContainsz%AutoGraphTestCase.assertGraphContains   s    	! 	  +r!   c                 <    | j                  || j                         y r   )assertInr#   )r%   r&   s     r    assertOpCreatedz!AutoGraphTestCase.assertOpCreated   s    MM'4>>*r!   c                 d    | j                  t        |      t        | j                        z         y r   )assertEmptysetr#   )r%   op_typess     r    assertOpsNotCreatedz%AutoGraphTestCase.assertOpsNotCreated   s"    S]S%889r!   c                 :    | j                  | j                         y r   )rc   r#   )r%   s    r    assertNoOpsCreatedz$AutoGraphTestCase.assertNoOpsCreated   s    T^^$r!   c                 b    | j                   j                  t        |   t	        |      f       y r   )r1   r$   r   assertEquallistr%   r   r   s     r    rj   zAutoGraphTestCase.assertEqual   s#    OOEG/d<=r!   c                 b    | j                   j                  t        |   t	        |      f       y r   )r1   r$   r   
assertLessrk   rl   s     r    rn   zAutoGraphTestCase.assertLess   s#    OOEG.T
;<r!   c                 b    | j                   j                  t        |   t	        |      f       y r   )r1   r$   r   assertGreaterEqualrk   rl   s     r    rp   z$AutoGraphTestCase.assertGreaterEqual   s#    OOEG6T
CDr!   c                 b    | j                   j                  t        |   t	        |      f       y r   )r1   r$   r   assertDictEqualrk   rl   s     r    rr   z!AutoGraphTestCase.assertDictEqual   s#    OOEG3T$Z@Ar!   c                     | j                   t        d       | j                  | | _         | j                   j                          y )Nz6cannot use more than one assertRaisesRuntime in a test)r2   
ValueErrorassertRaisesRegex	__enter__)r%   r   s     r    assertRaisesRuntimez%AutoGraphTestCase.assertRaisesRuntime   s<    ~~!OPP+T++T2DNNNr!   )NN)__name__
__module____qualname____doc__r   r,   r   rJ   rL   rQ   r^   ra   rf   rh   rj   rn   rp   rr   rw   __classcell__)r   s   @r    r	   r	      s]    8 BF# D 4,+:%>=EBr!   r	   )r{   rV   r<   r   r   tensorflow.python.eagerr   tensorflow.python.frameworkr   r   tensorflow.python.opsr   tensorflow.python.platformr   TestCaser	   rC   r!   r    <module>r      s7     	 
   0 4 + + +K Kr!   