diff --git a/tensorflow/core/profiler/utils/BUILD b/tensorflow/core/profiler/utils/BUILD
index f7c6d5496d5..1dc8836058d 100644
--- a/tensorflow/core/profiler/utils/BUILD
+++ b/tensorflow/core/profiler/utils/BUILD
@@ -240,8 +240,7 @@ cc_library(
         ":trace_utils",
         ":xplane_builder",
         ":xplane_visitor",
-        "//tensorflow/core:platform_base",
-        "//tensorflow/core/platform:types",
+        "//tensorflow/core:lib",
         "//tensorflow/core/profiler/protobuf:xplane_proto_cc",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/strings",
@@ -255,7 +254,7 @@ tf_cc_test(
         ":xplane_builder",
         ":xplane_utils",
         ":xplane_visitor",
-        "//tensorflow/core:platform_base",
+        "//tensorflow/core:lib",
         "//tensorflow/core:test",
         "//tensorflow/core:test_main",
         "//tensorflow/core/profiler/protobuf:xplane_proto_cc",
@@ -276,7 +275,7 @@ cc_library(
         ":xplane_builder",
         ":xplane_schema",
         ":xplane_utils",
-        "//tensorflow/core/platform:types",
+        "//tensorflow/core:lib",
         "//tensorflow/core/profiler/protobuf:xplane_proto_cc",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/strings",
diff --git a/tensorflow/core/profiler/utils/xplane_utils.cc b/tensorflow/core/profiler/utils/xplane_utils.cc
index 1389af2f16a..72a206b127c 100644
--- a/tensorflow/core/profiler/utils/xplane_utils.cc
+++ b/tensorflow/core/profiler/utils/xplane_utils.cc
@@ -24,6 +24,7 @@ limitations under the License.
 #include "absl/strings/string_view.h"
 #include "tensorflow/core/platform/env_time.h"
 #include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/protobuf.h"
 #include "tensorflow/core/platform/types.h"
 #include "tensorflow/core/profiler/protobuf/xplane.pb.h"
 #include "tensorflow/core/profiler/utils/timespan.h"
@@ -34,6 +35,24 @@ namespace tensorflow {
 namespace profiler {
 namespace {
 
+// Removes all elements from 'array' for which pred is true.
+template <typename T, typename Pred>
+void RemoveIf(protobuf::RepeatedPtrField<T>* array, Pred&& pred) {
+  T** const begin = array->mutable_data();
+  T** const end = begin + array->size();
+  T** write = begin;
+  while (write < end && !pred(*write)) ++write;
+  if (write == end) return;
+  // 'write' points to the first element to be removed.
+  for (T** scan = write + 1; scan < end; ++scan) {
+    if (!pred(*scan)) std::swap(*scan, *write++);
+  }
+  const int new_size = write - begin;
+  const int size = array->size();
+  DCHECK_GE(size, new_size);
+  array->DeleteSubrange(new_size, size - new_size);
+}
+
 // Creates a Timespan from an XEvent.
 // WARNING: This should only be used when comparing events from the same XLine.
 Timespan XEventTimespan(const XEvent& event) {
@@ -113,28 +132,18 @@ void AddOrUpdateStrStat(int64 metadata_id, absl::string_view value,
 }
 
 void RemovePlaneWithName(XSpace* space, absl::string_view name) {
-  auto* planes = space->mutable_planes();
-  planes->erase(
-      std::remove_if(planes->begin(), planes->end(),
-                     [&](const XPlane& plane) { return plane.name() == name; }),
-      planes->end());
+  RemoveIf(space->mutable_planes(),
+           [&](const XPlane* plane) { return plane->name() == name; });
 }
 
 void RemoveEmptyPlanes(XSpace* space) {
-  auto* planes = space->mutable_planes();
-  planes->erase(std::remove_if(planes->begin(), planes->end(),
-                               [&](const XPlane& plane) {
-                                 return plane.lines_size() == 0;
-                               }),
-                planes->end());
+  RemoveIf(space->mutable_planes(),
+           [&](const XPlane* plane) { return plane->lines().empty(); });
 }
 
 void RemoveEmptyLines(XPlane* plane) {
-  auto* lines = plane->mutable_lines();
-  lines->erase(std::remove_if(
-                   lines->begin(), lines->end(),
-                   [&](const XLine& line) { return line.events_size() == 0; }),
-               lines->end());
+  RemoveIf(plane->mutable_lines(),
+           [&](const XLine* line) { return line->events().empty(); });
 }
 
 bool XEventsComparator::operator()(const XEvent* a, const XEvent* b) const {
diff --git a/tensorflow/core/profiler/utils/xplane_utils_test.cc b/tensorflow/core/profiler/utils/xplane_utils_test.cc
index 04e06fcb05b..9aa4f9736d8 100644
--- a/tensorflow/core/profiler/utils/xplane_utils_test.cc
+++ b/tensorflow/core/profiler/utils/xplane_utils_test.cc
@@ -56,17 +56,28 @@ TEST(XPlaneUtilsTest, RemovePlaneWithName) {
   RemovePlaneWithName(&space, "non-exist");
   EXPECT_EQ(space.planes_size(), 0);
 
-  space.add_planes()->set_name("p1");
-  space.add_planes()->set_name("p2");
-  space.add_planes()->set_name("p3");
+  auto* p1 = space.add_planes();
+  p1->set_name("p1");
+  auto* p2 = space.add_planes();
+  p2->set_name("p2");
+  auto* p3 = space.add_planes();
+  p3->set_name("p3");
+
   RemovePlaneWithName(&space, "non-exist");
   EXPECT_EQ(space.planes_size(), 3);
+
   RemovePlaneWithName(&space, "p2");
   EXPECT_EQ(space.planes_size(), 2);
+  EXPECT_EQ(p1, FindPlaneWithName(space, "p1"));
+  EXPECT_EQ(p3, FindPlaneWithName(space, "p3"));
+
   RemovePlaneWithName(&space, "p1");
   EXPECT_EQ(space.planes_size(), 1);
+  EXPECT_EQ(p3, FindPlaneWithName(space, "p3"));
+
   RemovePlaneWithName(&space, "p1");
   EXPECT_EQ(space.planes_size(), 1);
+
   RemovePlaneWithName(&space, "p3");
   EXPECT_EQ(space.planes_size(), 0);
 }