Commit 5d6aec53 authored by A. Unique TensorFlower's avatar A. Unique TensorFlower Committed by TensorFlower Gardener
Browse files

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
parent d9d029f5
Loading
Loading
Loading
Loading
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please to comment