Track colocation context manager locations (file:line) and add colocation...
Track colocation context manager locations (file:line) and add colocation information support to error interpolation.
This CL add a new private property on ops: Operation._colocation_dict. This property will return a dictionary for which the keys are nodes with which this Operation is colocated, and for which the values are traceable_stack.TraceableObject instances. The TraceableObject instances record the location of the relevant colocation context manager but have the "obj" field set to None to prevent leaking private data.
For example, suppose file_a contained these lines:
file_a.py:
14: node_a = tf.constant(3, name='NODE_A')
15: with tf.colocate_with(node_a):
16: node_b = tf.constant(4, name='NODE_B')
Then a TraceableObject t_obj representing the colocation context manager would have these member values:
t_obj.obj -> None
t_obj.name = 'NODE_A'
t_obj.filename = 'file_a.py'
t_obj.lineno = 15
and node_b.op._colocation_dict would return the dictionary
{ 'NODE_A': t_obj }
PiperOrigin-RevId: 205035378
Loading
Please sign in to comment