Automated Code Change

PiperOrigin-RevId: 624422733
tensorflow · Apr 13, 2024 · 63756b7 · 63756b7
1 parent 0ecc83e
commit 63756b7
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/tensorflow/compiler/jit/pjrt_device_context.cc b/tensorflow/compiler/jit/pjrt_device_context.cc
@@ -93,7 +93,7 @@ void PjRtDeviceContext::CopyDeviceTensorToCPU(const Tensor* device_tensor,
                                               Device* device,
                                               Tensor* cpu_tensor,
                                               StatusCallback done) {
-  profiler::TraceMe traceme("PjRtDeviceContext::CopyDeviceTensorToCPU");
+  tsl::profiler::TraceMe traceme("PjRtDeviceContext::CopyDeviceTensorToCPU");
   if (device_tensor->NumElements() == 0) {
     VLOG(2) << "CopyDeviceTensorToCPU empty tensor";
     done(absl::OkStatus());
@@ -146,7 +146,7 @@ void PjRtDeviceContext::CopyCPUTensorToDevice(const Tensor* cpu_tensor,
                                               Tensor* device_tensor,
                                               StatusCallback done,
                                               bool sync_dst_compute) const {
-  profiler::TraceMe traceme("PjRtDeviceContext::CopyCPUTensorToDevice");
+  tsl::profiler::TraceMe traceme("PjRtDeviceContext::CopyCPUTensorToDevice");
   if (cpu_tensor->NumElements() == 0) {
     VLOG(2) << "CopyCPUTensorToDevice empty tensor";
     done(absl::OkStatus());
@@ -240,7 +240,7 @@ void PjRtDeviceToDeviceCopy(DeviceContext* send_dev_context,
                             AllocatorAttributes dst_alloc_attr,
                             const Tensor* input, Tensor* output,
                             int dev_to_dev_stream_index, StatusCallback done) {
-  profiler::TraceMe traceme("PjRtDevice_DeviceToDeviceCopy");
+  tsl::profiler::TraceMe traceme("PjRtDevice_DeviceToDeviceCopy");
   if (input->NumElements() == 0) {
     VLOG(2) << "PjRtDevice_DeviceToDeviceCopy empty tensor";
     done(absl::OkStatus());

diff --git a/tensorflow/compiler/jit/xla_device.cc b/tensorflow/compiler/jit/xla_device.cc
@@ -485,7 +485,8 @@ void XlaDevice::ComputeAsync(AsyncOpKernel* op_kernel, OpKernelContext* context,
 
 Status XlaDevice::Sync() {
   VLOG(1) << "XlaDevice::Sync";
-  profiler::TraceMe activity("XlaDevice::Sync", profiler::TraceMeLevel::kInfo);
+  tsl::profiler::TraceMe activity("XlaDevice::Sync",
+                                  tsl::profiler::TraceMeLevel::kInfo);
   std::shared_ptr<se::Stream> stream;
   {
     mutex_lock lock(mu_);