From 44dc41df400cdb72b42e02510c2b18917087be67 Mon Sep 17 00:00:00 2001
From: Yang Chen <yangchen@google.com>
Date: Fri, 21 Feb 2025 12:05:04 -0800
Subject: [PATCH] Cleanup: Fix includes.

PiperOrigin-RevId: 729611238
---
 tensorflow/core/kernels/data/experimental/BUILD                | 1 +
 .../core/kernels/data/experimental/rebatch_dataset_op.cc       | 3 ++-
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/tensorflow/core/kernels/data/experimental/BUILD b/tensorflow/core/kernels/data/experimental/BUILD
index 32c7cd028b34ae..b178ad4377d6dc 100644
--- a/tensorflow/core/kernels/data/experimental/BUILD
+++ b/tensorflow/core/kernels/data/experimental/BUILD
@@ -739,6 +739,7 @@ tf_kernel_library(
         "//tensorflow/core:lib_internal",
         "//tensorflow/core:protos_all_cc",
         "//tensorflow/core/data:name_utils",
+        "@com_google_absl//absl/strings",
     ],
 )
 
diff --git a/tensorflow/core/kernels/data/experimental/rebatch_dataset_op.cc b/tensorflow/core/kernels/data/experimental/rebatch_dataset_op.cc
index 711dcc0fe30437..8cb56e90b59724 100644
--- a/tensorflow/core/kernels/data/experimental/rebatch_dataset_op.cc
+++ b/tensorflow/core/kernels/data/experimental/rebatch_dataset_op.cc
@@ -15,6 +15,7 @@ limitations under the License.
 
 #include <optional>
 
+#include "absl/strings/str_join.h"
 #include "tensorflow/core/data/name_utils.h"
 #include "tensorflow/core/framework/dataset.h"
 #include "tensorflow/core/framework/tensor.h"
@@ -330,7 +331,7 @@ class RebatchDatasetV2Op : public UnaryDatasetOpKernel {
           output_types_(output_types),
           output_shapes_(output_shapes),
           traceme_metadata_(
-              {{"batch_sizes", absl::StrJoin(batch_sizes, ",")}}) {
+              {{"batch_sizes", absl::StrJoin(batch_sizes_, ",")}}) {
       input_->Ref();
     }