Automated Code Change

PiperOrigin-RevId: 1004523690

tensorflow/core/tpu/ops/BUILD

@@ -154,6 +154,7 @@ cc_library(
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/types:span",
         "@xla//xla:util",
+        "@xla//xla/tpu/protobuf:tpu_embedding_configuration_proto_cc",
     ],
 )
 
@@ -180,6 +181,8 @@ cc_library(
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:str_format",
         "@xla//xla:status_macros",
+        "@xla//xla/tpu/protobuf:optimization_parameters_proto_cc",
+        "@xla//xla/tpu/protobuf:tpu_embedding_configuration_proto_cc",
         "@xla//xla/tsl/platform:errors",
         "@xla//xla/tsl/platform:logging",
     ],

tensorflow/core/tpu/ops/tpu_embedding_ops.cc

@@ -25,6 +25,8 @@ limitations under the License.
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_format.h"
 #include "xla/status_macros.h"
+#include "xla/tpu/protobuf/optimization_parameters.pb.h"
+#include "xla/tpu/protobuf/tpu_embedding_configuration.pb.h"
 #include "xla/tsl/platform/errors.h"
 #include "xla/tsl/platform/logging.h"  // IWYU pragma: keep
 #include "tensorflow/core/framework/attr_value.pb.h"

tensorflow/core/tpu/ops/tpu_embedding_shape_util.cc

@@ -22,6 +22,7 @@ limitations under the License.
 #include "absl/status/statusor.h"
 #include "absl/strings/str_cat.h"
 #include "absl/types/span.h"
+#include "xla/tpu/protobuf/tpu_embedding_configuration.pb.h"
 #include "xla/tsl/platform/errors.h"
 #include "xla/tsl/platform/statusor.h"
 #include "xla/util.h"

tensorflow/core/tpu/ops/tpu_embedding_shape_util.h

@@ -21,6 +21,7 @@ limitations under the License.
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "absl/types/span.h"
+#include "xla/tpu/protobuf/tpu_embedding_configuration.pb.h"
 #include "tensorflow/core/framework/tensor_shape.pb.h"
 #include "tensorflow/core/platform/status.h"
 #include "tensorflow/core/protobuf/tpu/tpu_embedding_configuration.pb.h"