diff --git a/tensorflow/python/eager/BUILD b/tensorflow/python/eager/BUILD
index b4db76da05a..af5f3d16408 100644
--- a/tensorflow/python/eager/BUILD
+++ b/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",
],
)
diff --git a/tensorflow/python/eager/def_function.py b/tensorflow/python/eager/def_function.py
index 1b4e73329c6..d38870f2fe2 100644
--- a/tensorflow/python/eager/def_function.py
+++ b/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,
- tf_function_call="eager"):
+ with trace.Trace(self._name, 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())):
diff --git a/tensorflow/python/eager/function.py b/tensorflow/python/eager/function.py
index 857eb9cb656..c02318cb814 100644
--- a/tensorflow/python/eager/function.py
+++ b/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,
- tf_function_call="concrete"):
+ with trace.Trace(self._func_graph.name, 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:
diff --git a/tensorflow/python/eager/profiler_test.py b/tensorflow/python/eager/profiler_test.py
index f924f0f10f6..33a46bc0127 100644
--- a/tensorflow/python/eager/profiler_test.py
+++ b/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
diff --git a/tensorflow/python/profiler/BUILD b/tensorflow/python/profiler/BUILD
index c01a51314de..9996f5c9894 100644
--- a/tensorflow/python/profiler/BUILD
+++ b/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",
],
)
diff --git a/tensorflow/python/profiler/profiler_v2_test.py b/tensorflow/python/profiler/profiler_v2_test.py
index 50d29c0532f..42fbeba1e98 100644
--- a/tensorflow/python/profiler/profiler_v2_test.py
+++ b/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