"//tensorflow/core:lib",
        "//tensorflow/core:protos_all_cc",
        "//tensorflow/core/tpu:tpu_compile",
        "//tensorflow/core/tpu/kernels:tpu_compile_op_support",
        "//tensorflow/core/tpu/kernels:tpu_compile_proto_cc",
        "//tensorflow/core/tpu/kernels:tpu_util",
        "@com_google_absl//absl/container:flat_hash_set",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status",

        "//tensorflow/core:framework",
        "//tensorflow/core:lib",
        "//tensorflow/core/protobuf/tpu:compile_metadata_proto_cc",
        "//tensorflow/core/tpu:tpu_compile",
        "//tensorflow/core/tpu/kernels:tpu_compile_op_support",
        "@com_google_absl//absl/log",
        "@llvm-project//llvm:Support",
        "@llvm-project//mlir:IR",
        "@llvm-project//mlir:Pass",

  return absl::OkStatus();
}

// TODO(b/305211853): Unify the CPU/TPU/GPU Execution Ops and thus these two
// passes should merge together.
void RegisterTPULowerClusterToRuntimeOpsPassPipeline() {
  static mlir::PassPipelineRegistration<StandardPipelineOptions> pipeline(
      "tfrt-lower-cluster-to-runtime-ops-tpu",
      "Run all the passes involved after the clustering transformations from "

  pm.addPass(mlir::tf_executor::CreateTFExecutorTPUV1IslandOutliningPass());
}

// Run the TF XLA Bridge based on the input pipeline, which can be either TPU
// bridge pipeline or non TPU bridge pipeline.
tensorflow::Status RunTFXLABridge(
    ModuleOp module,
    llvm::function_ref<void(OpPassManager &pm)> pipeline_builder,
    llvm::StringRef module_name = llvm::StringRef(),

# ARM64 WHEEL
test:linux_arm64_wheel_test_filters --test_tag_filters=-no_oss,-no_aarch64,-oss_excluded,-oss_serial,-gpu,-tpu,-benchmark-test,-v1only,-no_oss_py38,-no_oss_py39,-no_oss_py310

  // For all other inputs that go from TPU op to TPU op, insert the
  // TPUOutput/Input pair.

  // Add TPUReplicatedInput/TPUReplicatedOutput pairs along each edge.
  // TODO(bfontain): Should be merged with the above loop.
  llvm::SetVector<Value> values_to_add_nodes;

  for (Operation* op : backward_pass_ops) {
    // TODO(bfontain): Should validate that all the TPU ops are in the same
    // replication region.

source ci/official/envs/linux_x86
TFCI_BAZEL_COMMON_ARGS="--repo_env=TF_PYTHON_VERSION=$TFCI_PYTHON_VERSION --config release_cpu_linux --config=tpu"
TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=linux_tpu
TFCI_BUILD_PIP_PACKAGE_ARGS="--repo_env=WHEEL_NAME=tensorflow_tpu"
TFCI_LIB_SUFFIX="-tpu-linux-x86_64"
TFCI_WHL_BAZEL_TEST_ENABLE=0
TFCI_WHL_IMPORT_TEST_ENABLE=0
TFCI_WHL_SIZE_LIMIT=580M

#include <variant>
#include <vector>

#include "absl/strings/string_view.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"

namespace tensorflow {
namespace tf2xla {
namespace internal {

// Fills in arg_shapes and metadata_proto with appropriate values based on the
// input mlir module.

  getOperation().walk([&](tf_device::ClusterFuncOp cluster_func) {
    cluster_funcs.push_back(cluster_func);
  });

  OpBuilder builder(&getContext());
  // Add resource reads for resource writes from TPU cluster where for such
  // resources the TPU cluster does not read from.
  for (tf_device::ClusterFuncOp cluster_func : cluster_funcs) {
    builder.setInsertionPoint(cluster_func);

    SmallVector<Value, 4> read_operands;

                                        const StringAttr &op_device) {
    // In TFRT TPU, tensor transfer is handled specifically by D2H and
    // H2D transfer kernels. So fold the tf.Identity op if:
    // * the identity op is placed on TPU, and
    // * the arg to the identity op is produced by a TPUExecuteOp.
    if (op_device && op_device.getValue().contains("TPU")) {
      return true;
    }

    Operation *def_op = arg.getDefiningOp();