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();

#include "xla/xla.pb.h"
#include "xla/xla_data.pb.h"
#include "tensorflow/core/framework/tensor_shape.pb.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"

namespace mlir {
namespace TFTPU {

// Populates a TPUCompileMetadataProto from attributes of a
// `tf_device::ClusterFuncOp`. If any necessary attributes are missing from the

    compatible_with = get_compatible_with_portable(),  # copybara: comment
    visibility = [
        # copybara:uncomment "//learning/brain/experimental/tfrt:__subpackages__",
        # copybara:uncomment "//learning/brain/tfrt/tpu/compiler:__subpackages__",
        "//tensorflow/compiler/mlir/tfrt:__subpackages__",
        "//tensorflow/core/runtime_fallback:internal",
        "//tensorflow/core/tfrt/saved_model:friends",

#include <string>

#include "mlir/Support/LogicalResult.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/utils/tpu_rewrite_device_util.h"

namespace mlir {
namespace TFTPU {

// For each TPU cluster in `module`, walk over all ops inside the cluster
// and reachable in the call graph from the cluster.
// For each op walked, `callback` is applied to the op, the root cluster, and

  return llvm::isa<TF::CastOp, TF::IdentityOp>(op);
}

// Adds outside compilation attributes to unary ops such as Identity/Cast ops
// at the head of TPU computation that is used only by other outside compiled
// ops. Identity ops and Cast ops is commonly added to the start of TPU
// computation. Adding/expanding outside compilation attributes to these ops
// will ensure that head outside compiled ops are correctly located and moved to
// host.

// _tpu_replicate or _xla_compile_device_type=TPU. The function library will be
// skipped if nullptr is provided.

bool IsSupportedByReplicatedBridge(
    const Graph& graph, const FunctionLibraryDefinition* function_library);

// Check if an MLIR module has any ops with _tpu_replicate or
// _xla_compile_device_type=TPU.
bool IsSupportedByReplicatedBridge(mlir::ModuleOp module);

// RUN: tf-opt %s -tf-replicated-clustering-bridge-v2 -tfrt-lower-cluster-to-runtime-ops-tpu 2>&1 | FileCheck %s

// This test verifies that the tail extraction is not terminated prematurely
// due to the outside compilation attribute could be removed in
// canonicalization of Reshape ops.

// Reshape should not be executed on TPU as all are marked by outside
// compilation. And there should be no host-device communication.
// CHECK: tf._TPUCompileMlir

    %control = tf_executor.island wraps "tf.TPUReplicateMetadata"() {_tpu_replicate = "cluster", allow_soft_placement = false, computation_shape = [], device = "", device_assignment = [], host_compute_core...

        // Handle replicated TPU case.
        islands_to_outline.push_back(island_op);
        break;
      }
      auto device_type =
          op.getAttrOfType<StringAttr>(TF::kCompileDeviceTypeAttr);
      if (device_type && device_type.getValue() == TF::kTpuDevice &&
          !op.hasAttrOfType<StringAttr>(TF::kReplicationInfoAttr)) {
        // Handle single-core TPU case (no `TPUReplicateMetadataOp`).

// RUN: tf-opt %s -tf-replicated-clustering-bridge-v2 -tfrt-lower-cluster-to-runtime-ops-tpu 2>&1 | FileCheck %s

// This test verifies that the tail extraction is not terminated prematurely
// in handling tf.If op which would end up with excessive host-device
// communication.

// In this test, all ops other than tf.Rank are marked with outside_compilation
// . So the TPU program should contain tf.Rank op and there should be no
// host-device communication.