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

1 files changed, 1 insertions, 0 deletions

diff --git a/tensorflow/python/BUILD b/tensorflow/python/BUILD
index a362dee97d..c33a579ad2 100644
--- a/tensorflow/python/BUILD
+++ b/tensorflow/python/BUILD
@@ -1044,6 +1044,7 @@ py_test(
         ":client_testlib",
         ":constant_op",
         ":error_interpolation",
+        ":traceable_stack",
     ],
 )