o
    i e_<                     @   s   d Z ddlmZ ddlmZ ddlmZ dd Zdd Z	d	d
 Z
dd Zdd Zdd ZG dd deZG dd deZdd ZG dd deZdd ZdS )z=Classes and methods for processing debugger-decorated graphs.    )	graph_pb2)op_def_registry)
tf_loggingc                 C   sN   d| v r#|  ds#| d| d }t| | dd d }||fS | dfS )a  Get the node name from a string that can be node or tensor name.

  Args:
    name: An input node name (e.g., "node_a") or tensor name (e.g.,
      "node_a:0"), as a str.

  Returns:
    1) The node name, as a str. If the input name is a tensor name, i.e.,
      consists of a colon, the final colon and the following output slot
      will be stripped.
    2) If the input name is a tensor name, the output slot, as an int. If
      the input name is not a tensor name, None.
  :N   )endswithrfindint)name	node_nameoutput_slot r   W/var/www/myenv/lib/python3.10/site-packages/tensorflow/python/debug/lib/debug_graphs.pyparse_node_or_tensor_name   s
   r   c                 C   s   t | \}}|S Nr   )element_namer   _r   r   r   get_node_name-   s   r   c                 C   s   t | \}}|dur|S dS )a  Get the output slot number from the name of a graph element.

  If element_name is a node name without output slot at the end, 0 will be
  assumed.

  Args:
    element_name: (`str`) name of the graph element in question.

  Returns:
    (`int`) output slot number.
  Nr   r   )r   r   r   r   r   r   get_output_slot2   s   r   c                 C   
   |  dS )a=  Determine whether a node name is that of a debug Copy node.

  Such nodes are inserted by TensorFlow core upon request in
  RunOptions.debug_options.debug_tensor_watch_opts.

  Args:
    node_name: Name of the node.

  Returns:
    A bool indicating whether the input argument is the name of a debug Copy
    node.
  __copy_
startswithr   r   r   r   is_copy_nodeB   s   
r   c                 C   r   )a/  Determine whether a node name is that of a debug node.

  Such nodes are inserted by TensorFlow core upon request in
  RunOptions.debug_options.debug_tensor_watch_opts.

  Args:
    node_name: Name of the node.

  Returns:
    A bool indicating whether the input argument is the name of a debug node.
  __dbg_r   r   r   r   r   is_debug_nodeR   s   
r   c                 C   s   d}| }| |std|  |t|d }|ddk r$td|  ||dd d }|d|d }t||dd d }|d|d }|ddkr[td	|  |d|d }t||dd d }||||fS )
a~  Parse the name of a debug node.

  Args:
    node_name: Name of the debug node.

  Returns:
    1. Name of the watched node, as a str.
    2. Output slot index of the watched tensor, as an int.
    3. Index of the debug node, as an int.
    4. Name of the debug op, as a str, e.g, "DebugIdentity".

  Raises:
    ValueError: If the input node name is not a valid debug node name.
  r   z'Invalid prefix in debug node name: '%s'Nr      zInvalid debug node name: '%s'r   r   z,Invalid tensor name in debug node name: '%s')r   
ValueErrorlencountrindexr	   index)r   prefixr
   debug_opdebug_op_indexwatched_node_namewatched_output_slotr   r   r   parse_debug_node_namea   s    
r)   c                   @   s   e Zd ZdS )GraphTracingReachedDestinationN)__name__
__module____qualname__r   r   r   r   r*      s    r*   c                   @   s6   e Zd ZdZ		dddZdd Zdd Zd	d
 ZdS )DFSGraphTracerz,Graph input tracer using depth-first search.Nc                 C   s.   || _ || _g | _g | _d| _g | _|| _dS )a  Constructor of _DFSGraphTracer.

    Args:
      input_lists: A list of dicts. Each dict is an adjacency (input) map from
        the recipient node name as the key and the list of input node names
        as the value.
      skip_node_names: Optional: a list of node names to skip tracing.
      destination_node_name: Optional: destination node name. If not `None`, it
        should be the name of a destination not as a str and the graph tracing
        will raise GraphTracingReachedDestination as soon as the node has been
        reached.

    Raises:
      GraphTracingReachedDestination: if stop_at_node_name is not None and
        the specified node is reached.
    r   N)_input_lists_skip_node_names_inputs_visited_nodes_depth_count_depth_list_destination_node_name)selfinput_listsskip_node_namesdestination_node_namer   r   r   __init__   s   
zDFSGraphTracer.__init__c                 C   s   |  j d7  _ t|}|| jkrt || jv rdS || jv r!dS | j| | jD ](}||vr1q*|| D ]}t|| jv r?q5| j| | j	| j  | 
| q5q*|  j d8  _ dS )a  Trace inputs.

    Args:
      graph_element_name: Name of the node or an output tensor of the node, as a
        str.

    Raises:
      GraphTracingReachedDestination: if destination_node_name of this tracer
        object is not None and the specified node is reached.
    r   N)r3   r   r5   r*   r0   r2   appendr/   r1   r4   trace)r6   graph_element_namer   
input_listinpr   r   r   r<      s(   



zDFSGraphTracer.tracec                 C      | j S r   )r1   r6   r   r   r   inputs      zDFSGraphTracer.inputsc                 C   r@   r   )r4   rA   r   r   r   
depth_list   rC   zDFSGraphTracer.depth_list)NN)r+   r,   r-   __doc__r:   r<   rB   rD   r   r   r   r   r.      s    
$r.   c                 C   s6   d}| j D ]
}|jr|j} nq|du rtd |S )z,Infer device name from a partition GraphDef.NzsFailed to infer device name from partition GraphDef: none of the nodes of the GraphDef has a non-empty device name.)nodedeviceloggingwarn)	graph_defdevice_namerF   r   r   r   _infer_device_name   s   
rL   c                   @   s   e Zd ZdZ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edd Zedd Zedd Zedd Zedd Zedd  Zed!d" Zed#d$ Zed%d& Zed'd( Zed)d* ZdS ),
DebugGraphz&Represents a debugger-decorated graph.Nc                 C   s   || _ d | _i | _i | _i | _i | _i | _i | _i | _i | _	g | _
i | _|| _| js/t|| _|jD ]}| | q2|   |   | |   |   d S r   )_debug_graph_def_non_debug_graph_def_node_attributes_node_inputs_node_reversed_ref_inputs_node_ctrl_inputs_node_recipients_node_ctrl_recipients_node_devices_node_op_types_copy_send_nodes	_ref_args_device_namerL   rF   _process_debug_graph_node%_prune_non_control_edges_of_debug_ops!_prune_control_edges_of_debug_ops*_prune_nodes_from_input_and_recipient_maps_get_copy_nodes_populate_recipient_maps)r6   debug_graph_defrK   rF   r   r   r   r:      s*   

zDebugGraph.__init__c                 C   s<  t |jrdS |j| jv rtd| j|jf |j| j|j< g | j|j< g | j|j< g | j|j< g | j	|j< |j| j
vrCt | j
|j< | j
|j |jrO|jn| j |j| j|j< | || j|j< |jD ]5}t|r}|jdksv|jdkr}| j|j |dr|dd }| j|j | qf| j|j | qfdS )zProcess a node from the debug GraphDef.

    Args:
      node: (NodeDef) A partition-graph node to be processed.

    Raises:
      ValueError: If duplicate node names are encountered.
    Nz&Duplicate node name on device %s: '%s'_Send_Retval^r   )r   r
   rQ   r   rZ   attrrP   rS   rT   rU   rV   setaddrG   oprW   _get_ref_argsrY   inputr   rX   r;   r   )r6   rF   r?   cinpr   r   r   r[     s4   
	


z$DebugGraph._process_debug_graph_nodec                 C   s`   t |j}|du rg S g }t|jD ]\}}|jr-|dkr!|jnd|j|f }|| q|S )zDetermine whether an input of an op is ref-type.

    Args:
      node: A `NodeDef`.

    Returns:
      A list of the arg names (as strs) that are ref-type.
    Nr   z%s:%d)r   getrh   	enumerate
output_argis_refr
   r;   )r6   rF   op_defref_argsirn   arg_namer   r   r   ri   3  s   	
zDebugGraph._get_ref_argsc                 C   s&   g }| j D ]}t|r|| q|S )z(Find all Copy nodes in the loaded graph.)rQ   r   r;   )r6   
copy_nodesrF   r   r   r   r_   G  s   

zDebugGraph._get_copy_nodesc                 C   sJ   | j D ]}| j | }t|D ]\}}t|r!| j | d }|||< qqdS )zPrune (non-control) edges related to debug ops.

    Prune the Copy ops and associated _Send ops inserted by the debugger out
    from the non-control inputs and output recipients map. Replace the inputs
    and recipients with original ones.
    r   N)rQ   rm   r   )r6   rF   rB   rr   r?   orig_inpr   r   r   r\   O  s   

z0DebugGraph._prune_non_control_edges_of_debug_opsc                 C   sN   | j D ]!}| j | }g }|D ]}t|r|| q|D ]}|| qqdS )z-Prune control edges related to the debug ops.N)rS   r   r;   remove)r6   rF   ctrl_inputsdebug_op_inputsctrl_inpdebug_op_inpr   r   r   r]   `  s   


z,DebugGraph._prune_control_edges_of_debug_opsc                 C   s   | j D ]9}| j | }|D ]/}t|}|| jvrg | j|< | j| | || jv r;|| jvr3g | j|< | j| | qq| jD ]$}| j| }|D ]}|| jv rQqI|| jvr[g | j|< | j| | qIq@dS )zPopulate the map from node name to recipient(s) of its output(s).

    This method also populates the input map based on reversed ref edges.
    N)	rQ   r   rT   r;   rY   rR   rS   rX   rU   )r6   rF   rB   r?   rw   ry   r   r   r   r`   k  s.   











z#DebugGraph._populate_recipient_mapsc                 C   s.   |D ]}| j |= | j|= | j|= | j|= qdS )zPrune nodes out of input and recipient maps.

    Args:
      nodes_to_prune: (`list` of `str`) Names of the nodes to be pruned.
    N)rQ   rS   rT   rU   )r6   nodes_to_prunerF   r   r   r   r^     s   
z5DebugGraph._prune_nodes_from_input_and_recipient_mapsc                 C   s   | j rdS t | _ | jjD ]>}t|jst|jrq| j j }|	| |j
dd= | j|j D ]}|j
| q2| j|j D ]
}|j
d|  qAqdS )zReconstruct non-debug GraphDef.

    Non-debug GraphDef means the original GraphDef without the Copy* and Debug
    nodes inserted by the debugger.
    Nrd   )rO   r   GraphDefrN   rF   r   r
   r   rg   CopyFromrj   rQ   r;   rS   )r6   rF   new_noder?   ry   r   r   r    _reconstruct_non_debug_graph_def  s   

z+DebugGraph._reconstruct_non_debug_graph_defc                 C   r@   r   )rZ   rA   r   r   r   rK        zDebugGraph.device_namec                 C   r@   )z The debugger-decorated GraphDef.)rN   rA   r   r   r   ra     s   zDebugGraph.debug_graph_defc                 C   s   |    | jS )zFThe GraphDef without the Copy* and Debug* nodes added by the debugger.)r   rO   rA   r   r   r   non_debug_graph_def  s   zDebugGraph.non_debug_graph_defc                 C   r@   r   )rV   rA   r   r   r   node_devices  r   zDebugGraph.node_devicesc                 C   r@   r   )rW   rA   r   r   r   node_op_types  r   zDebugGraph.node_op_typesc                 C   r@   r   )rP   rA   r   r   r   node_attributes  r   zDebugGraph.node_attributesc                 C   r@   r   )rQ   rA   r   r   r   node_inputs  r   zDebugGraph.node_inputsc                 C   r@   r   )rS   rA   r   r   r   node_ctrl_inputs  r   zDebugGraph.node_ctrl_inputsc                 C   r@   r   )rR   rA   r   r   r   node_reversed_ref_inputs  r   z#DebugGraph.node_reversed_ref_inputsc                 C   r@   r   )rT   rA   r   r   r   node_recipients  r   zDebugGraph.node_recipientsc                 C   r@   r   )rU   rA   r   r   r   node_ctrl_recipients  r   zDebugGraph.node_ctrl_recipientsr   )r+   r,   r-   rE   r:   r[   ri   r_   r\   r]   r`   r^   r   propertyrK   ra   r   r   r   r   r   r   r   r   r   r   r   r   r   rM      sB    
+









rM   c                 C   s
   t | jS )a  Reconstruct original (non-debugger-decorated) partition GraphDef.

  This method strips the input `tf.compat.v1.GraphDef` of the Copy* and
  Debug*-type nodes inserted by the debugger.

  The reconstructed partition graph is identical to the original (i.e.,
    non-debugger-decorated) partition graph except in the following respects:
      1) The exact names of the runtime-inserted internal nodes may differ.
         These include _Send, _Recv, _HostSend, _HostRecv, _Retval ops.
      2) As a consequence of 1, the nodes that receive input directly from such
         send- and recv-type ops will have different input names.
      3) The parallel_iteration attribute of while-loop Enter ops are set to 1.

  Args:
    debug_graph_def: The debugger-decorated `tf.compat.v1.GraphDef`, with the
      debugger-inserted Copy* and Debug* nodes.

  Returns:
    The reconstructed `tf.compat.v1.GraphDef` stripped of the debugger-inserted
    nodes.
  )rM   r   )ra   r   r   r   reconstruct_non_debug_graph_def  s   
r   N)rE   tensorflow.core.frameworkr   tensorflow.python.frameworkr   tensorflow.python.platformr   rH   r   r   r   r   r   r)   	Exceptionr*   objectr.   rL   rM   r   r   r   r   r   <module>   s    )M u