Automated Code Change

PiperOrigin-RevId: 629359187
tensorflow · May 1, 2024 · ba448e0 · ba448e0
1 parent 40573b4
commit ba448e0
Show file tree

Hide file tree

Showing 16 changed files with 47 additions and 52 deletions.
diff --git a/tensorflow/compiler/mlir/lite/quantization/ir/QuantOpsBase.td b/tensorflow/compiler/mlir/lite/quantization/ir/QuantOpsBase.td
@@ -27,7 +27,6 @@ include "mlir/IR/OpBase.td"
 def QuantizationFork_Dialect : Dialect {
  let name = "quantfork";
  let cppNamespace = "::mlir::quantfork";
- let usePropertiesForAttributes = 0;
 }
 
 #endif // QUANT_FORK_BASE
diff --git a/tensorflow/compiler/mlir/lite/quantization/tests/import_quant_stats.mlir b/tensorflow/compiler/mlir/lite/quantization/tests/import_quant_stats.mlir
@@ -18,8 +18,8 @@ func.func @import_stats_name(%arg0: tensor<4xf32>, %cst: tensor<i32>) -> (tensor
  func.return %0#0, %0#1 : tensor<2xf32>, tensor<2xf32>
 
 // CHECK-NEXT: %[[split:.*]]:2 = "tfl.split"
-// CHECK-NEXT: %[[stats1:.*]] = "quantfork.stats"(%[[split]]#0) {layerStats = dense<[-1.000000e+00, 1.000000e+00]>
-// CHECK-NEXT: %[[stats2:.*]] = "quantfork.stats"(%[[split]]#1) {layerStats = dense<[-1.000000e+00, 1.000000e+00]>
+// CHECK-NEXT: %[[stats1:.*]] = "quantfork.stats"(%[[split]]#0) <{layerStats = dense<[-1.000000e+00, 1.000000e+00]>
+// CHECK-NEXT: %[[stats2:.*]] = "quantfork.stats"(%[[split]]#1) <{layerStats = dense<[-1.000000e+00, 1.000000e+00]>
 // CHECK-NEXT: return %[[stats1]], %[[stats2]] : tensor<2xf32>, tensor<2xf32>
 }
 
@@ -30,7 +30,7 @@ func.func @import_stats_name_port(%arg0: tensor<4xf32>, %cst: tensor<i32>) -> (t
  func.return %0#0, %0#1 : tensor<2xf32>, tensor<2xf32>
 
 // CHECK-NEXT: %[[split:.*]]:2 = "tfl.split"
-// CHECK-NEXT: %[[stats1:.*]] = "quantfork.stats"(%[[split]]#0) {layerStats = dense<[-2.000000e+00, 2.000000e+00]>
+// CHECK-NEXT: %[[stats1:.*]] = "quantfork.stats"(%[[split]]#0) <{layerStats = dense<[-2.000000e+00, 2.000000e+00]>
 // CHECK-NEXT: return %[[stats1]], %[[split]]#1 : tensor<2xf32>, tensor<2xf32>
 }
 
@@ -41,7 +41,7 @@ func.func @import_stats_name_regex(%arg0: tensor<4xf32>, %cst: tensor<i32>) -> (
  func.return %0#0, %0#1 : tensor<2xf32>, tensor<2xf32>
 
 // CHECK-NEXT: %[[split:.*]]:2 = "tfl.split"
-// CHECK-NEXT: %[[stats1:.*]] = "quantfork.stats"(%[[split]]#0) {layerStats = dense<[-3.000000e+00, 3.000000e+00]>
-// CHECK-NEXT: %[[stats2:.*]] = "quantfork.stats"(%[[split]]#1) {layerStats = dense<[-3.000000e+00, 3.000000e+00]>
+// CHECK-NEXT: %[[stats1:.*]] = "quantfork.stats"(%[[split]]#0) <{layerStats = dense<[-3.000000e+00, 3.000000e+00]>
+// CHECK-NEXT: %[[stats2:.*]] = "quantfork.stats"(%[[split]]#1) <{layerStats = dense<[-3.000000e+00, 3.000000e+00]>
 // CHECK-NEXT: return %[[stats1]], %[[stats2]] : tensor<2xf32>, tensor<2xf32>
 }
diff --git a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/importer_test_min_max.cc b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/importer_test_min_max.cc
@@ -100,14 +100,14 @@ std::optional<std::unique_ptr<tflite::ModelT>> InjectStatsToFullyConnected(
 
  // CHECK-LABEL: func @main(%arg0: tensor<40x37xf32>, %arg1: tensor<40x37xf32>)
  // CHECK-SAME: -> tensor<40x40xf32>
- // CHECK: %[[stat:.*]] = "quantfork.stats"(%arg0) {layerStats = dense<
- // CHECK-SAME: [-1.000000e+00, 1.000000e+00]> : tensor<2xf32>}
+ // CHECK: %[[stat:.*]] = "quantfork.stats"(%arg0) <{layerStats = dense
+ // CHECK-SAME: <[-1.000000e+00, 1.000000e+00]> : tensor<2xf32>}>
  // CHECK-SAME: : (tensor<40x37xf32>) -> tensor<40x37xf32>
  // CHECK-NEXT: %[[cst:.*]] = "tfl.pseudo_const"() {value = dense<
  // CHECK-SAME: 1.000000e+00> : tensor<40xf32>} : () -> tensor<40xf32>
  // CHECK-NEXT: %[[fc:.*]]:2 = "tfl.fully_connected"(%[[stat]], %arg1,
  // CHECK-NEXT: %[[stat1:.*]] = "quantfork.stats"(%[[fc]]#0)
- // CHECK-SAME: {axis = 1 : i64,
+ // CHECK-SAME: <{axis = 1 : i64,
  // CHECK-SAME: axisStats = dense<{{\[}}[-0.000000e+00, 0.000000e+00],
  // CHECK-SAME: [-1.000000e+00, 1.000000e+00],
  // CHECK-SAME: [-2.000000e+00, 2.000000e+00]

diff --git a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/lstm.json b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/lstm.json
@@ -1,8 +1,8 @@
 // RUN: json_to_flatbuffer %p/test_schema.fbs %s | flatbuffer_translate --tflite-flatbuffer-to-mlir -o - | FileCheck %s
 // RUN: json_to_flatbuffer %p/test_schema.fbs %s | flatbuffer_translate --tflite-flatbuffer-to-mlir -o - | flatbuffer_translate -mlir-to-tflite-flatbuffer - -o - | flatbuffer_to_string - | FileCheck --check-prefix=RoundTrip %s
 
-// CHECK-DAG: %[[input_18:.*]] = "quantfork.stats"({{.*}}) {layerStats = dense<[-8.000000e-01, 1.600000e+00]> : tensor<2xf32>} : (tensor<1x4xf32>) -> tensor<1x4xf32>
-// CHECK-DAG: %[[input_19:.*]] = "quantfork.stats"({{.*}}) {layerStats = dense<[-2.000000e+00, 4.000000e+00]> : tensor<2xf32>} : (tensor<1x2xf32>) -> tensor<1x2xf32>
+// CHECK-DAG: %[[input_18:.*]] = "quantfork.stats"({{.*}}) <{layerStats = dense<[-8.000000e-01, 1.600000e+00]> : tensor<2xf32>}> : (tensor<1x4xf32>) -> tensor<1x4xf32>
+// CHECK-DAG: %[[input_19:.*]] = "quantfork.stats"({{.*}}) <{layerStats = dense<[-2.000000e+00, 4.000000e+00]> : tensor<2xf32>}> : (tensor<1x2xf32>) -> tensor<1x2xf32>
 
 // CHECK: "tfl.unidirectional_sequence_lstm"({{.*}}, %[[input_18]], %[[input_19]], %{{[0-9]+}}, %{{[0-9]+}}, %{{[0-9]+}}, %{{[0-9]+}})
 // CHECK-SAME: effective_hidden_scale_intermediate = tensor<*x!quant.calibrated<f32<-5.000000e-01:5.000000e-01>>>

diff --git a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/quant_stats.mlir b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/quant_stats.mlir
@@ -3,7 +3,7 @@
 
 func.func @main(%arg0: tensor<1x512x672x8xf32>) -> tensor<1x512x672x8xf32> {
 // CHECK-LABEL: @main
-// CHECK: %[[RES0:.*]] = "quantfork.stats"(%arg0) {layerStats = dense<[0.000000e+00, 2.550000e+02]> : tensor<2xf32>} : (tensor<1x512x672x8xf32>) -> tensor<1x512x672x8xf32>
+// CHECK: %[[RES0:.*]] = "quantfork.stats"(%arg0) <{layerStats = dense<[0.000000e+00, 2.550000e+02]> : tensor<2xf32>}> : (tensor<1x512x672x8xf32>) -> tensor<1x512x672x8xf32>
 
  %0 = "quantfork.stats"(%arg0) {layerStats = dense<[0.000000e+00, 2.550000e+02]> : tensor<2xf32>} : (tensor<1x512x672x8xf32>) -> tensor<1x512x672x8xf32>
  func.return %0 : tensor<1x512x672x8xf32>

diff --git a/tensorflow/compiler/mlir/quantization/common/ir/QuantOpsBase.td b/tensorflow/compiler/mlir/quantization/common/ir/QuantOpsBase.td
@@ -27,7 +27,6 @@ include "mlir/IR/OpBase.td"
 def Quant_Dialect : Dialect {
  let name = "quantization";
  let cppNamespace = "::mlir::quant::ir";
- let usePropertiesForAttributes = 0;
 }
 
 #endif // QUANTIZATION_BASE
diff --git a/...iler/mlir/quantization/tensorflow/tests/convert_custom_aggregation_op_to_quant_stats.mlir b/...iler/mlir/quantization/tensorflow/tests/convert_custom_aggregation_op_to_quant_stats.mlir
@@ -6,7 +6,7 @@ func.func @customAggregator(%arg0: tensor<8x8x8x8xf32>) -> (tensor<8x8x8x8xf32>,
  func.return %0#0, %1#0 : tensor<8x8x8x8xf32>, tensor<8x8x8x8xf32>
 }
 // CHECK: func @customAggregator
-// CHECK-NEXT: %[[stats:.*]] = "quantfork.stats"(%arg0) {layerStats = dense<[-1.000000e-01, 2.000000e-01]> : tensor<2xf32>} : (tensor<8x8x8x8xf32>) -> tensor<8x8x8x8xf32>
+// CHECK-NEXT: %[[stats:.*]] = "quantfork.stats"(%arg0) <{layerStats = dense<[-1.000000e-01, 2.000000e-01]> : tensor<2xf32>}> : (tensor<8x8x8x8xf32>) -> tensor<8x8x8x8xf32>
 // CHECK-NEXT: return %[[stats]], %arg0
 
 func.func @doNotHandleNoMinMaxCases(%arg0: tensor<8x8x8x8xf32>) -> (tensor<8x8x8x8xf32>, tensor<8x8x8x8xf32>, tensor<8x8x8x8xf32>) {

diff --git a/tensorflow/compiler/mlir/tf2xla/api/v2/BUILD b/tensorflow/compiler/mlir/tf2xla/api/v2/BUILD
@@ -22,77 +22,66 @@ cc_library(
  ],
  deps = [
  ":device_type_proto_cc",
- "//tensorflow/compiler/jit:flags_headers",
- "//tensorflow/compiler/jit:shape_inference",
- "//tensorflow/compiler/mlir/tensorflow",
  "//tensorflow/compiler/mlir/tensorflow:dump_mlir_util",
- "//tensorflow/compiler/mlir/tensorflow:error_util",
- "//tensorflow/compiler/mlir/tensorflow:export_graphdef",
- "//tensorflow/compiler/mlir/tensorflow:mlir_roundtrip_flags",
- "//tensorflow/compiler/mlir/tensorflow:serialize_mlir_module_utils",
- "//tensorflow/compiler/mlir/tensorflow:translate_utils",
- "//tensorflow/compiler/mlir/tensorflow/transforms:set_tpu_infeed_layout",
- "//tensorflow/compiler/mlir/tensorflow/transforms:tensorflow_passes",
  "//tensorflow/compiler/mlir/tf2xla:compile_mlir_util",
  "//tensorflow/compiler/mlir/tf2xla/api/v1:compile_tf_graph",
  "//tensorflow/compiler/mlir/tf2xla/internal:compilation_timer",
  "//tensorflow/compiler/mlir/tf2xla/internal:legalize_tf_mlir",
  "//tensorflow/compiler/mlir/tf2xla/internal:legalize_tf_to_hlo",
  "//tensorflow/compiler/tf2xla:layout_util",
- "//tensorflow/compiler/tf2xla:tf2xla_util",
  "//tensorflow/compiler/tf2xla:xla_helpers",
  "//tensorflow/core:core_cpu_base",
  "//tensorflow/core:framework",
  "//tensorflow/core:lib",
- "//tensorflow/core/platform:statusor",
- "//tensorflow/core/tpu:tpu_compile",
+ "//tensorflow/core/protobuf:for_core_protos_cc",
+ "//tensorflow/core/protobuf/tpu:compile_metadata_proto_cc",
  "//tensorflow/core/tpu/kernels:tpu_compile_op_support",
  "//tensorflow/core/tpu/kernels:tpu_compile_proto_cc",
- "//tensorflow/core/tpu/kernels:tpu_util_hdrs",
  "@com_google_absl//absl/log",
  "@com_google_absl//absl/status",
+ "@com_google_absl//absl/status:statusor",
  "@com_google_absl//absl/strings",
  "@com_google_absl//absl/types:variant",
  "@llvm-project//llvm:Support",
  "@llvm-project//mlir:IR",
  "@llvm-project//mlir:Pass",
  "@local_tsl//tsl/platform:error_logging",
- "@local_tsl//tsl/platform:status",
  "@local_tsl//tsl/platform:statusor",
+ "@local_xla//xla:shape_util",
  "@local_xla//xla:xla_proto_cc",
  "@local_xla//xla/client:compile_only_client",
  "@local_xla//xla/hlo/ir:hlo",
- "@local_xla//xla/mlir_hlo:hlo_dialect_registration",
  "@local_xla//xla/pjrt:compile_options_proto_cc",
- "@stablehlo//:register",
  ],
 )
 
 tf_cc_test(
  name = "legalize_tf_test",
  srcs = ["legalize_tf_test.cc"],
  deps = [
- ":device_type_proto_cc",
  ":legalize_tf",
  "//tensorflow/compiler/jit",
- "//tensorflow/compiler/mlir/tensorflow",
  "//tensorflow/compiler/mlir/tf2xla/internal:test_matchers",
  "//tensorflow/compiler/mlir/tf2xla/internal/utils:test_metadata_config",
  "//tensorflow/compiler/tf2xla:xla_compiler",
  "//tensorflow/compiler/tf2xla:xla_helpers",
  "//tensorflow/core:framework",
  "//tensorflow/core:lib",
+ "//tensorflow/core:protos_all_cc",
  "//tensorflow/core:test",
  "//tensorflow/core:test_main",
- "//tensorflow/core/lib/monitoring:cell_reader",
  "//tensorflow/core/protobuf:for_core_protos_cc",
  "//tensorflow/core/protobuf/tpu:compile_metadata_proto_cc",
  "//tensorflow/core/tpu/kernels:tpu_compile_op_support",
+ "@com_google_absl//absl/status:statusor",
  "@com_google_absl//absl/strings:str_format",
  "@com_google_googletest//:gtest",
+ "@llvm-project//mlir:Pass",
  "@local_tsl//tsl/lib/monitoring:test_utils",
  "@local_tsl//tsl/platform:statusor",
+ "@local_xla//xla:shape_util",
  "@local_xla//xla/client:client_library",
+ "@local_xla//xla/stream_executor:platform",
  "@local_xla//xla/stream_executor:platform_manager",
  ],
 )
@@ -132,6 +121,7 @@ cc_library(
  "//tensorflow/core/platform:errors",
  "//tensorflow/core/platform:stacktrace",
  "//tensorflow/core/platform:status",
+ "//tensorflow/core/protobuf:for_core_protos_cc",
  "@com_google_absl//absl/log",
  "@com_google_absl//absl/status",
  "@llvm-project//llvm:Support",
@@ -158,7 +148,6 @@ tf_cc_test(
  "//tensorflow/compiler/mlir/tensorflow",
  "//tensorflow/compiler/mlir/tensorflow:attribute_utils",
  "//tensorflow/compiler/mlir/tensorflow:tf_dialect_lib",
- "//tensorflow/core/lib/monitoring:cell_reader",
  "//tensorflow/core/platform:resource_loader",
  "@com_google_absl//absl/status",
  "@com_google_absl//absl/strings",
@@ -167,7 +156,6 @@ tf_cc_test(
  "@llvm-project//mlir:IR",
  "@llvm-project//mlir:Parser",
  "@local_tsl//tsl/lib/core:status_test_util",
- "@local_tsl//tsl/platform:status",
  ],
 )
 
@@ -183,13 +171,13 @@ cc_library(
  ],
  deps = [
  "//tensorflow/compiler/jit:flags_headers",
- "//tensorflow/compiler/mlir/tensorflow:bridge_logger",
  "//tensorflow/compiler/mlir/tensorflow:dump_mlir_util",
  "//tensorflow/compiler/mlir/tensorflow/transforms:verify_no_outside_compilation_markers_pass",
  "//tensorflow/compiler/mlir/tf2xla/internal:logging_hooks",
  "//tensorflow/core:framework",
  "//tensorflow/core/platform:error_payloads",
  "//tensorflow/core/platform:status",
+ "//tensorflow/core/protobuf:for_core_protos_cc",
  "@com_google_absl//absl/log",
  "@com_google_absl//absl/status",
  "@llvm-project//llvm:Support",
@@ -200,8 +188,6 @@ cc_library(
  "@llvm-project//mlir:Transforms",
  "@local_tsl//tsl/lib/monitoring:counter",
  "@local_tsl//tsl/platform:error_logging",
- "@local_tsl//tsl/platform:errors",
- "@local_tsl//tsl/platform:status",
  ],
 )
 
@@ -215,15 +201,12 @@ tf_cc_test(
  deps = [
  ":tf_dialect_to_executor",
  "//tensorflow/compiler/mlir:register_common_dialects",
- "//tensorflow/core/lib/monitoring:cell_reader",
  "//tensorflow/core/platform:resource_loader",
  "@com_google_absl//absl/status",
  "@com_google_absl//absl/strings",
  "@com_google_googletest//:gtest_main",
  "@llvm-project//mlir:IR",
  "@llvm-project//mlir:Parser",
  "@local_tsl//tsl/lib/core:status_test_util",
- "@local_tsl//tsl/lib/monitoring:test_utils",
- "@local_tsl//tsl/platform:status",
  ],
 )
diff --git a/tensorflow/compiler/mlir/tf2xla/api/v2/cluster_tf.cc b/tensorflow/compiler/mlir/tf2xla/api/v2/cluster_tf.cc
@@ -39,6 +39,7 @@ limitations under the License.
 #include "tensorflow/core/platform/errors.h"
 #include "tensorflow/core/platform/stacktrace.h"
 #include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/protobuf/core_platform_payloads.pb.h"
 #include "tensorflow/core/util/debug_data_dumper.h"
 #include "tsl/platform/error_logging.h"
 #include "tsl/platform/errors.h"

diff --git a/tensorflow/compiler/mlir/tf2xla/api/v2/cluster_tf_test.cc b/tensorflow/compiler/mlir/tf2xla/api/v2/cluster_tf_test.cc
@@ -30,12 +30,9 @@ limitations under the License.
 #include "mlir/Parser/Parser.h" // from @llvm-project
 #include "tensorflow/compiler/mlir/register_common_dialects.h"
 #include "tensorflow/compiler/mlir/tensorflow/ir/tf_device.h"
-#include "tensorflow/compiler/mlir/tensorflow/ir/tf_executor.h"
 #include "tensorflow/compiler/mlir/tensorflow/utils/attribute_utils.h"
-#include "tensorflow/core/lib/monitoring/cell_reader.h"
 #include "tensorflow/core/platform/resource_loader.h"
 #include "tsl/lib/core/status_test_util.h"
-#include "tsl/platform/status.h"
 
 namespace tensorflow {
 namespace tf2xla {

diff --git a/tensorflow/compiler/mlir/tf2xla/api/v2/legalize_tf.cc b/tensorflow/compiler/mlir/tf2xla/api/v2/legalize_tf.cc
@@ -26,6 +26,8 @@ limitations under the License.
 #include "absl/strings/str_cat.h"
 #include "absl/types/variant.h"
 #include "llvm/ADT/ScopeExit.h"
+#include "llvm/ADT/StringRef.h"
+#include "mlir/Pass/Pass.h" // from @llvm-project
 #include "tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.h"
 #include "tensorflow/compiler/mlir/tf2xla/api/v1/compile_mlir_util.h"
 #include "tensorflow/compiler/mlir/tf2xla/api/v1/compile_tf_graph.h"
@@ -34,11 +36,16 @@ limitations under the License.
 #include "tensorflow/compiler/mlir/tf2xla/internal/legalize_tf_to_hlo.h"
 #include "tensorflow/compiler/tf2xla/layout_util.h"
 #include "tensorflow/compiler/tf2xla/xla_helpers.h"
+#include "xla/client/compile_only_client.h"
 #include "xla/hlo/ir/hlo_module.h"
+#include "xla/shape.h"
 #include "xla/xla.pb.h"
 #include "tensorflow/core/framework/metrics.h"
+#include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/graph/graph.h"
 #include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/protobuf/config.pb.h"
+#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"
 #include "tensorflow/core/tpu/kernels/tpu_compile_op_support.h"
 #include "tensorflow/core/util/debug_data_dumper.h"
 #include "tensorflow/core/util/dump_graph.h"

diff --git a/tensorflow/compiler/mlir/tf2xla/api/v2/legalize_tf.h b/tensorflow/compiler/mlir/tf2xla/api/v2/legalize_tf.h
@@ -21,13 +21,17 @@ limitations under the License.
 #include <vector>
 
 #include "absl/status/status.h"
+#include "absl/status/statusor.h"
 #include "absl/types/variant.h"
+#include "llvm/ADT/StringRef.h"
 #include "mlir/IR/BuiltinOps.h" // from @llvm-project
 #include "mlir/Pass/Pass.h" // from @llvm-project
 #include "tensorflow/compiler/mlir/tf2xla/api/v2/device_type.pb.h"
+#include "tensorflow/compiler/tf2xla/layout_util.h"
 #include "tensorflow/compiler/tf2xla/xla_helpers.h"
 #include "xla/client/compile_only_client.h"
 #include "xla/pjrt/compile_options.pb.h"
+#include "xla/shape.h"
 #include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/tpu/kernels/tpu_compile.pb.h"
 #include "tensorflow/core/tpu/kernels/tpu_compile_op_support.h"

diff --git a/tensorflow/compiler/mlir/tf2xla/api/v2/legalize_tf_test.cc b/tensorflow/compiler/mlir/tf2xla/api/v2/legalize_tf_test.cc
@@ -22,14 +22,22 @@ limitations under the License.
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
+#include "absl/status/statusor.h"
 #include "absl/strings/str_format.h"
+#include "mlir/Pass/Pass.h" // from @llvm-project
 #include "tensorflow/compiler/mlir/tf2xla/internal/test_matchers.h"
 #include "tensorflow/compiler/mlir/tf2xla/internal/utils/test_metadata_config.h"
 #include "tensorflow/compiler/tf2xla/xla_compiler.h"
-#include "tensorflow/compiler/tf2xla/xla_helpers.h"
 #include "xla/client/client_library.h"
+#include "xla/shape.h"
+#include "xla/stream_executor/platform.h"
 #include "xla/stream_executor/platform_manager.h"
-#include "tensorflow/core/lib/monitoring/cell_reader.h"
+#include "tensorflow/core/framework/attr_value.pb.h"
+#include "tensorflow/core/framework/function.h"
+#include "tensorflow/core/framework/function.pb.h"
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/lib/monitoring/test_utils.h"
 #include "tensorflow/core/platform/env.h"
 #include "tensorflow/core/platform/test.h"

diff --git a/tensorflow/compiler/mlir/tf2xla/api/v2/tf_dialect_to_executor.cc b/tensorflow/compiler/mlir/tf2xla/api/v2/tf_dialect_to_executor.cc
@@ -29,15 +29,14 @@ limitations under the License.
 #include "mlir/Transforms/Passes.h" // from @llvm-project
 #include "tensorflow/compiler/jit/flags.h"
 #include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"
-#include "tensorflow/compiler/mlir/tensorflow/utils/data_dumper_logger_config.h"
 #include "tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.h"
 #include "tensorflow/compiler/mlir/tf2xla/internal/logging_hooks.h"
 #include "tensorflow/core/platform/error_payloads.h"
 #include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/protobuf/core_platform_payloads.pb.h"
 #include "tensorflow/core/util/debug_data_dumper.h"
 #include "tsl/lib/monitoring/counter.h"
 #include "tsl/platform/error_logging.h"
-#include "tsl/platform/status.h"
 
 namespace tensorflow {
 namespace tf2xla {

diff --git a/tensorflow/compiler/mlir/tf2xla/api/v2/tf_dialect_to_executor_test.cc b/tensorflow/compiler/mlir/tf2xla/api/v2/tf_dialect_to_executor_test.cc
@@ -27,10 +27,8 @@ limitations under the License.
 #include "mlir/IR/OwningOpRef.h" // from @llvm-project
 #include "mlir/Parser/Parser.h" // from @llvm-project
 #include "tensorflow/compiler/mlir/register_common_dialects.h"
-#include "tensorflow/core/lib/monitoring/cell_reader.h"
 #include "tensorflow/core/platform/resource_loader.h"
 #include "tsl/lib/core/status_test_util.h"
-#include "tsl/platform/status.h"
 
 namespace tensorflow {
 namespace tf2xla {

diff --git a/tensorflow/core/ops/ops.pbtxt b/tensorflow/core/ops/ops.pbtxt
@@ -1,4 +1,4 @@
-go/debugstr 
+go/debugstr  
 op {
  name: "Abort"
  attr {