
    BVhFD                     \   d Z ddlZddlmZ ddlmZ  eg d      Zd Zd Z	d]dZ
 ej                  d	d
      d        Z ej                  dd
      d        Z ej                  dd
      d        Z ej                  dd
      d        Z ej                  dd
      d        Z ej                  dd
      d        Z ej                  dd
      d        Z ej                  dd
      d        Z ej                  dd
      d        Zd]dZ ej                  dd
       ej                  dd
      d               Z ej                  d d
      d!        Z ej                  d"d
      d#        Z ej                  d$d
      d%        Z ej                  d&d
      d'        Z ej                  d(d
      d)        Z ej                  d*d
      d+        Z ej                  d,d
      d-        Z ej                  d.d
      d/        Z ej                  d0d
      d1        Z ej                  d2d
      d3        Z  ej                  d4d
      d5        Z! ej                  d6d
      d7        Z" ej                  d8d
      d9        Z# ej                  d:d
      d;        Z$d^d<Z% ej                  d=d
      d>        Z& ej                  d?d
      d@        Z' ej                  dAd
      dB        Z( ej                  dCd
      dD        Z) ej                  dEd
      dF        Z*dG Z+dH Z, ej                  dId
      dJ        Z- ej                  dKd
      dL        Z. ej                  dMd
      dN        Z/ ej                  dOd
      dP        Z0 ej                  dQd
      dR        Z1 ej                  dSd
      dT        Z2 ej                  dUd
      dV        Z3 ej                  dWd
      dX        Z4 ej                  dYd
       ej                  dZd
       ej                  d[d
      d\                      Z5y)_z=Register flops statistics for various TensorFlow operations.
    N)
graph_util)ops)*
ReciprocalSquareRsqrtLogNeg	AssignSub	AssignAddL2LossSoftmaxAddSubMulRealDivMaximumMinimumPow	RsqrtGradGreaterEqualGreater	LessEqualLessEqualNotEqualSquaredDifferenceAddV2MeanSumArgMaxArgMinBiasAddGradAvgPoolMaxPoolAvgPoolGradMaxPoolGradConv2DBackpropInputConv2DBackpropFilterAddNMatMulConv2DDepthwiseConv2dNativeBiasAdd
Dilation2Dc                 2    ~ ~t        j                  dd      S )zReturns zero flops.flopsr   )r   OpStatsgraphnodes     b/home/dcms/DCMS/lib/python3.12/site-packages/tensorflow/python/profiler/internal/flops_registry.py_zero_flopsr6   ,   s    T	Wa	      c                 "    d}| D ]  }||z  }	 |S )z(Computes product of element of the list.    )lstresultitems      r5   _list_productr>   2   s$    & d
dNF	-r7   c                     t        j                  | |j                  d         }|j                          t	        j
                  d|j                         |z        S )z5Common code which compute flops for unary operations.r   r0   r   tensor_shape_from_node_def_nameinputassert_is_fully_definedr   r1   num_elements)r3   r4   ops_per_elementin_shapes       r5   _unary_op_flopsrG   >   sG    77tzz!}M(
""$	Wh335G	HHr7   r   r0   c                     t        | |      S )z'Compute flops for Reciprocal operation.rG   r2   s     r5   _reciprocal_flopsrJ   E        
	%%r7   r   c                     t        | |      S )z#Compute flops for Square operation.rI   r2   s     r5   _square_flopsrM   K   rK   r7   r   c                     t        | |d      S )z"Compute flops for Rsqrt operation.   rE   rI   r2   s     r5   _rsqrt_flopsrQ   Q   s     
a	88r7   r   c                     t        | |      S )z Compute flops for Log operation.rI   r2   s     r5   
_log_flopsrS   X   rK   r7   r	   c                     t        | |      S )z Compute flops for Neg operation.rI   r2   s     r5   
_neg_flopsrU   ^   rK   r7   r
   c                     t        | |      S )z&Compute flops for AssignSub operation.rI   r2   s     r5   _assign_sub_flopsrW   d   rK   r7   r   c                     t        | |      S )z&Compute flops for AssignAdd operation.rI   r2   s     r5   _assign_add_flopsrY   j   rK   r7   r   c                     t        j                  | |j                  d         }|j                          t	        j
                  d|j                         dz  dz
        S )z#Compute flops for L2Loss operation.r   r0      r9   r@   r3   r4   rF   s      r5   _l2_loss_flopsr]   p   sP     77tzz!}M(
""$ 
Wh3359A=	>>r7   r   c                     t        | |d      S )z$Compute flops for Softmax operation.   rP   rI   r2   s     r5   _softmax_flopsr`   z   s     
a	88r7   c                     t        j                  | |j                        }|j                          t	        j
                  d|j                         |z        S )z6Common code which compute flops for binary operations.r0   )r   rA   namerC   r   r1   rD   )r3   r4   rE   	out_shapes       r5   _binary_per_element_op_flopsrd      sC    88		J)##%	Wi446H	IIr7   r   r   c                     t        | |      S )z Compute flops for Add operation.rd   r2   s     r5   
_add_flopsrg      s     
&eT	22r7   r   c                     t        | |      S )z Compute flops for Sub operation.rf   r2   s     r5   
_sub_flopsri           
&eT	22r7   r   c                     t        | |      S )z Compute flops for Mul operation.rf   r2   s     r5   
_mul_flopsrl      rj   r7   r   c                     t        | |      S )z$Compute flops for RealDiv operation.rf   r2   s     r5   _real_div_flopsrn      rj   r7   r   c                     t        | |      S )z$Compute flops for Maximum operation.rf   r2   s     r5   _maximum_flopsrp      rj   r7   r   c                     t        | |      S )z$Compute flops for Minimum operation.rf   r2   s     r5   _minimum_flopsrr      rj   r7   r   c                     t        | |      S )z Compute flops for Pow operation.rf   r2   s     r5   
_pow_flopsrt      rj   r7   r   c                     t        | |d      S )z&Compute flops for RsqrtGrad operation.   rP   rf   r2   s     r5   _rsqrt_grad_flopsrw           
&eT1	EEr7   r   c                     t        | |      S )z)Compute flops for GreaterEqual operation.rf   r2   s     r5   _greater_equal_flopsrz      rj   r7   r   c                     t        | |      S )z$Compute flops for Greater operation.rf   r2   s     r5   _greater_flopsr|      rj   r7   r   c                     t        | |      S )z&Compute flops for LessEqual operation.rf   r2   s     r5   _less_equal_flopsr~      rj   r7   r   c                     t        | |      S )z!Compute flops for Less operation.rf   r2   s     r5   _less_flopsr      rj   r7   r   c                     t        | |      S )z"Compute flops for Equal operation.rf   r2   s     r5   _equal_flopsr      rj   r7   r   c                     t        | |      S )z%Compute flops for NotEqual operation.rf   r2   s     r5   _not_equal_flopsr      rj   r7   r   c                     t        | |d      S )z.Compute flops for SquaredDifference operation.rO   rP   rf   r2   s     r5   _squared_difference_flopsr      rx   r7   c                 H   t        j                  | |j                  d         }|j                          t        j                  | |j                        }|j                          |j                         |z  |j                         ||z
  z  z   }t        j                  d|      S )z9Common code which compute flops for reduction operations.r   r0   )r   rA   rB   rC   rb   rD   r   r1   )r3   r4   reduce_flopsfinalize_flopsrF   rc   	num_flopss          r5   _reduction_op_flopsr      s    77tzz!}M(
""$88		J)##%$$&5'')^l-JKL)	Wi	((r7   r   c                      t        | |dd      S )z!Compute flops for Mean operation.r9   r   r   r   r2   s     r5   _mean_flopsr           
UDq	KKr7   r   c                      t        | |dd      S )z Compute flops for Sum operation.r9   r   r   r   r2   s     r5   
_sum_flopsr     r   r7   r    c                      t        | |dd      S )z#Compute flops for ArgMax operation.r9   r   r   r   r2   s     r5   _arg_max_flopsr   	  r   r7   r!   c                      t        | |dd      S )z#Compute flops for ArgMin operation.r9   r   r   r   r2   s     r5   _arg_min_flopsr     r   r7   r"   c                      t        | |dd      S )z(Compute flops for BiasAddGrad operation.r9   r   r   r   r2   s     r5   _bias_add_grad_flopsr     s    
 
UDq	KKr7   c                 R    | j                   d   j                  dk7  rt        d      y)z>Verifies data format for pooling and convolutional operations.data_formats   NHWCz3Only NHWC format is supported in flops computationsN)attrs
ValueError)r4   s    r5   _verify_conv_data_formatr   $  s,     
YY}7*
J
KK +r7   c                 4   t        |       t        j                  | |j                        }|j	                          t        |j                  d   j
                  j                        }t        |      }t        j                  d||j                         z        S )z7Common code which compute flops for pooling operations.ksizer0   )r   r   rA   rb   rC   listr   ir>   r   r1   rD   )r3   r4   rc   kernel_shapekernel_areas        r5   _pool_flopsr   +  sx     4  88		J)##%dii(--//0,l++	WkI,B,B,DD	EEr7   r#   c                     t        | |      S )z$Compute flops for AvgPool operation.r   r2   s     r5   _avg_pool_flopsr   C       
UD	!!r7   r$   c                     t        | |      S )z$Compute flops for MaxPool operation.r   r2   s     r5   _max_pool_flopsr   I  r   r7   r%   c                 @   t        |       t        j                  | |j                  d         }|j	                          t        |j                  d   j
                  j                        }t        |      }t        j                  d||j                         z  dz        S )z(Compute flops for AvgPoolGrad operation.r9   r   r0   rO   )r   r   rA   rB   rC   r   r   r   r>   r   r1   rD   )r3   r4   out_backprop_shaper   r   s        r5   _avg_pool_grad_flopsr   O  s     4 !AA%BF**Q-Q,,.dii(--//0,l++ 
W #5#B#B#DDqH
J Jr7   r&   c                 `   t        |       t        |j                  d   j                  j                        }t	        |      }t        j                  | |j                  d         }|j                          ||j                         z  }t        j                  d||j                         z         S )z(Compute flops for MaxPoolGrad operation.r   r9   r0   )r   r   r   r   r>   r   rA   rB   rC   rD   r   r1   )r3   r4   r   r   orig_out_shapemax_pool_opss         r5   _max_pool_grad_flopsr   `  s     4   dii(--//0,l++==e>BjjmM.((*~::<<,	Wl^-H-H-JJ	KKr7   r'   c                    t        |       t        j                  | |j                        }|j	                          t        j                  | |j
                  d         }|j	                          t        |j                  d   j                  j                        }|d   |d   z  }t        j                  dd|j                         z  |j                         z  |j                  d   j                  |z  z        S )z0Compute flops for Conv2DBackpropInput operation.r9   stridesrO   r0   )r   r   rA   rb   rC   rB   r   r   r   r   r1   rD   dimsvalue)r3   r4   rc   r   strides_shapestrides_products         r5   _conv_2d_backprop_input_flopsr   |  s     4 88		J)##%;;E<@JJqMK,&&(tyy+00223-!!$}Q'77/	W)0022$1134"+11OCE
G Gr7   r(   c                    t        |       t        j                  | |j                  d         }|j	                          t        j                  | |j
                        }|j	                          t        |j                  d   j                  j                        }|d   |d   z  }t        j                  dd|j                         z  |j                         z  |j                  d   j                  |z  z        S )z1Compute flops for Conv2DBackpropFilter operation.r   r   r9   rO   r0   r   )r   r   rA   rB   rC   rb   r   r   r   r   r1   rD   r   r   )r3   r4   image_shaper   r   r   s         r5   _conv_2d_backprop_filter_flopsr     s     4 ::5$**Q-P+%%';;E499M,&&(tyy+00223-!!$}Q'77/	W+2244$1134$))"-33oEG
I Ir7   r)   c                    |j                   st        | |      S t        j                  | |j                   d         }|j	                          t        j                  d|j                         t        |j                         dz
  z        S )z!Compute flops for AddN operation.r   r0   r9   )	rB   r6   r   rA   rC   r   r1   rD   lenr\   s      r5   _add_n_flopsr     sj     
ud##77tzz!}M(
""$	Wh335TZZ19LM	NNr7   r*   c                    |j                   d   j                  }t        j                  | |j                  d         }|j                          |rt        |d         }nt        |d         }t        j                  | |j                        }|j                          t        j                  |j                               }t        j                  d||z  dz        S )z3Calculates the compute resources needed for MatMul.transpose_ar   r9   r0   rO   r   br   rA   rB   rC   intrb   npprodas_listr   r1   r3   r4   r   a_shapekoutput_shapeoutput_counts          r5   _calc_mat_mul_flopsr     s     		-(**+66udjjmL'	!!#GAJAGAJA;;E499M,&&(--/0,	Wq</!3	55r7   BatchMatMulBatchMatMulV2BatchMatMulV3c                    |j                   d   j                  }t        j                  | |j                  d         }|j                          |rt        |d         }nt        |d         }t        j                  | |j                        }|j                          t        j                  |j                               }t        j                  d||z  dz        S )z8Calculates the compute resources needed for BatchMatMul.r   r   r   r0   rO   r   r   s          r5   _calc_batch_mat_mul_flopsr     s    
 		-(**+66udjjmL'	!!#GBKAGBKA;;E499M,&&(--/0,	Wq</!3	55r7   )r9   )r9   r   )6__doc__numpyr   tensorflow.python.frameworkr   r   setIMPLEMENTED_OPSr6   r>   rG   RegisterStatisticsrJ   rM   rQ   rS   rU   rW   rY   r]   r`   rd   rg   ri   rl   rn   rp   rr   rt   rw   rz   r|   r~   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r:   r7   r5   <module>r      s    2 +   (!I g.& /&
 '*& +&
 )9 *9 w'& (&
 w'& (&
 W-& .&
 W-& .&
 '*? +? 	7+9 ,9J w')3 * (3
 w'3 (3
 w'3 (3
 	7+3 ,3
 	7+3 ,3
 	7+3 ,3
 w'3 (3
 W-F .F
 03 13
 	7+3 ,3
 W-3 .3
 (3 )3
 )3 *3
 
G,3 -3
 +W5F 6F) (L )L w'L (L '*L +L '*L +L w/L 0LLF0 	7+" ,"
 	7+" ,"
 w/J 0J  w/L 0L6 -w7G 8G> .8I 9I4 (O )O '*6 +6 w/116 2 2 06r7   