Use trace.Trace instead of traceme.TraceMe in python

PiperOrigin-RevId: 316220088
Change-Id: I92f28c7afa1f12136e7668f280fa93cecff837b9

tensorflow/python/eager/BUILD

@@ -275,7 +275,7 @@ cuda_py_test(
         ":test",
         "//tensorflow/core/profiler/protobuf:trace_events_proto_py",
         "//tensorflow/python:constant_op",
-        "//tensorflow/python/profiler:traceme",
+        "//tensorflow/python/profiler:trace",
     ],
 )
 
@@ -769,7 +769,7 @@ py_library(
         "//tensorflow/python:util",
         "//tensorflow/python:variable_scope",
         "//tensorflow/python:while_v2",  # TODO(b/118513001): Imported via control_flow_ops; remove.
-        "//tensorflow/python/profiler:traceme",
+        "//tensorflow/python/profiler:trace",
         "//tensorflow/python/training/tracking:base",
     ],
 )

tensorflow/python/eager/def_function.py

@@ -39,7 +39,7 @@ from tensorflow.python.ops import control_flow_util
 from tensorflow.python.ops import math_ops
 from tensorflow.python.ops import resource_variable_ops
 from tensorflow.python.platform import tf_logging as logging
-from tensorflow.python.profiler import traceme
+from tensorflow.python.profiler import trace
 from tensorflow.python.training.tracking import base as trackable
 from tensorflow.python.util import deprecation
 from tensorflow.python.util import nest
@@ -757,12 +757,11 @@ class Function(object):
   def __call__(self, *args, **kwds):
     """Calls the graph function and warn too frequent tracings."""
     if RUN_FUNCTIONS_EAGERLY:
-      with traceme.TraceMe(self._name,
+      with trace.Trace(self._name, tf_function_call="eager"):
-                           tf_function_call="eager"):
         return self._python_function(*args, **kwds)
 
     tracing_count = self._get_tracing_count()
-    with traceme.TraceMe(self._name) as tm:
+    with trace.Trace(self._name) as tm:
       if self._experimental_compile and (
           not control_flow_util.GraphOrParentsInXlaContext(
               ops.get_default_graph())):

tensorflow/python/eager/function.py

@@ -65,7 +65,7 @@ from tensorflow.python.ops import gradients_util
 from tensorflow.python.ops import resource_variable_ops
 
 from tensorflow.python.platform import tf_logging as logging
-from tensorflow.python.profiler import traceme
+from tensorflow.python.profiler import trace
 from tensorflow.python.util import compat
 from tensorflow.python.util import function_utils
 from tensorflow.python.util import lazy_loader
@@ -1650,8 +1650,7 @@ class ConcreteFunction(object):
 
   def _call_impl(self, args, kwargs, cancellation_manager=None):
     """See `__call__` for details."""
-    with traceme.TraceMe(self._func_graph.name,
+    with trace.Trace(self._func_graph.name, tf_function_call="concrete"):
-                         tf_function_call="concrete"):
       # Construct the list of input tensors: check if the structured signature
       # applies first; and if not, then use the flat signature.
       if self._function_spec is not None:

tensorflow/python/eager/profiler_test.py

@@ -27,14 +27,14 @@ from tensorflow.python.framework import config
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import test_util
 from tensorflow.python.platform import gfile
-from tensorflow.python.profiler import traceme
+from tensorflow.python.profiler import trace
 
 
 class ProfilerTest(test_util.TensorFlowTestCase):
 
   def test_profile(self):
     profiler.start()
-    with traceme.TraceMe('three_times_five'):
+    with trace.Trace('three_times_five'):
       three = constant_op.constant(3)
       five = constant_op.constant(5)
       product = three * five

tensorflow/python/profiler/BUILD

@@ -67,7 +67,7 @@ cuda_py_test(
         ":profiler_v2",
         "//tensorflow/python:constant_op",
         "//tensorflow/python/eager:test",
-        "//tensorflow/python/profiler:traceme",
+        "//tensorflow/python/profiler:trace",
     ],
 )
 

tensorflow/python/profiler/profiler_v2_test.py

@@ -27,7 +27,7 @@ from tensorflow.python.framework import errors
 from tensorflow.python.framework import test_util
 from tensorflow.python.platform import gfile
 from tensorflow.python.profiler import profiler_v2 as profiler
-from tensorflow.python.profiler import traceme
+from tensorflow.python.profiler import trace
 
 
 class ProfilerTest(test_util.TensorFlowTestCase):
@@ -55,7 +55,7 @@ class ProfilerTest(test_util.TensorFlowTestCase):
   def test_save_profile(self):
     logdir = self.get_temp_dir()
     profiler.start(logdir)
-    with traceme.TraceMe('three_times_five'):
+    with trace.Trace('three_times_five'):
       three = constant_op.constant(3)
       five = constant_op.constant(5)
       product = three * five
@@ -91,7 +91,7 @@ class ProfilerTest(test_util.TensorFlowTestCase):
     options = profiler.ProfilerOptions(
         host_tracer_level=3, python_tracer_level=1)
     profiler.start(logdir, options)
-    with traceme.TraceMe('three_times_five'):
+    with trace.Trace('three_times_five'):
       three = constant_op.constant(3)
       five = constant_op.constant(5)
       product = three * five
@@ -106,7 +106,7 @@ class ProfilerTest(test_util.TensorFlowTestCase):
     options = profiler.ProfilerOptions(
         host_tracer_level=3, python_tracer_level=1)
     with profiler.Profile(logdir, options):
-      with traceme.TraceMe('three_times_five'):
+      with trace.Trace('three_times_five'):
         three = constant_op.constant(3)
         five = constant_op.constant(5)
         product = three * five