// expected-warning @+1 {{TF2XLA TPU bridge input check: cluster op = tf.opA with cluster = cluster has successor as non cluster op tf.opB}}
    %out2, %c2 = tf_executor.island wraps "tf.opB"(%out) : (tensor<i32>) -> tensor<i32>
    // expected-error @+1 {{tf.TPUReplicatedOutput' op TF2XLA TPU bridge input check: non-cluster op = tf.opB has invalid successor op = tf.TPUReplicatedOutput}}

#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"
#include "tsl/platform/errors.h"

namespace tensorflow {
namespace tf2xla {
namespace internal {
namespace {

constexpr char kEntryFuncName[] = "main";

absl::Status SetupArguments(mlir::ModuleOp module,
                            std::vector<TensorShape>& arg_shapes,
                            tpu::TPUCompileMetadataProto& metadata_proto) {

// RUN: tf-tfrt-opt -tfrt-lower-cluster-to-runtime-ops-tpu -split-input-file -verify-diagnostics %s | FileCheck %s

module attributes {tf.versions = {producer = 888 : i32}, tf.devices = ["/job:worker/replica:0/task:0/device:CPU:0", "/job:worker/replica:0/task:0/device:TPU_SYSTEM:0", "/job:worker/replica:0/task:0/device:TPU:0"]} {

  // CHECK-LABEL: @converts_cluster
  func.func @converts_cluster() {
    // CHECK: %0:2 = "tf_device.launch"() <{{.*}}> ({

    arg->set_dtype(dtype);
    // TODO(lyandy): Support other arg kinds.
    if (dtype == tensorflow::DT_RESOURCE)
      arg->set_kind(tensorflow::tpu::TPUCompileMetadataProto::Arg::VARIABLE);
    else
      arg->set_kind(tensorflow::tpu::TPUCompileMetadataProto::Arg::PARAMETER);

    // Populate argument shapes.
    *arg->mutable_shape() = tensorflow::TensorShapeProto();
    if (auto ranked_tensor_type =

// TPU-annotated operations and intended to preserve backward compatibility with
// TFv1.
struct TpuV1BridgeExecutorIslandCoarsening
    : public impl::TpuV1BridgeExecutorIslandCoarseningPassBase<
          TpuV1BridgeExecutorIslandCoarsening> {
  void runOnOperation() override;
};

// Returns name of TPU cluster, if op belongs to a TPU cluster. Otherwise,
// returns `std::nullopt`.

  mlir::tf_device::ClusterFuncOp cluster_func_op = cluster_func_ops.front();

  tensorflow::tpu::TPUCompileMetadataProto compile_metadata;

  ASSERT_TRUE(mlir::succeeded(SetMetadataProtoFromClusterFuncOp(
      cluster_func_op,
      /*num_replicas=*/1, /*num_cores_per_replica=*/1, {}, &compile_metadata)));

  tensorflow::tpu::TPUCompileMetadataProto expected_compile_metadata;
  ASSERT_TRUE(tsl::protobuf::TextFormat::ParseFromString(

#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"
#include "tensorflow/core/tpu/kernels/tpu_compile_op_support.h"
#include "tsl/platform/statusor.h"

namespace tensorflow {
namespace tf2xla {
namespace internal {

using ::tensorflow::monitoring::testing::CellReader;
using tpu::MlirToHloArgs;
using tpu::ShardingAndIndex;
using tpu::TPUCompileMetadataProto;

namespace tensorflow {
namespace tf2xla {
namespace v1 {
namespace {

using ::tensorflow::monitoring::testing::CellReader;
using ::tensorflow::tpu::FunctionToHloArgs;
using ::tensorflow::tpu::MlirToHloArgs;
using ::tensorflow::tpu::ShardingAndIndex;
using ::tsl::monitoring::testing::Histogram;

static constexpr char kCompilationTimeStreamzName[] =
    "/tensorflow/core/tf2xla/api/v1/phase2_compilation_time";

    // CHECK-NEXT:       %[[FUNC_CALL:[0-9]*]] = func.call @__reduce_func_1(%[[ARG_5]], %[[GET_VALUE]])
    // CHECK-SAME:       _xla_compile_device_type = "TPU"
    // CHECK-SAME:       device = "/job:localhost/replica:0/task:0/device:TPU:1"
    // CHECK:            "tf.Yield"(%[[FUNC_CALL]])
    // CHECK:            "tf.Yield"(%[[ARG_5]])

#include "tensorflow/core/protobuf/error_codes.pb.h"
#include "tensorflow/core/tfrt/common/create_pjrt_client_util.h"
#include "tensorflow/core/tfrt/common/pjrt_util.h"
#include "tensorflow/core/tpu/tpu_defs.h"

namespace tensorflow {
namespace {
using XlaDeviceCompiler =
    DeviceCompiler<xla::LocalExecutable, xla::LocalClient>;
using PjRtDeviceCompiler =