if (failed(TF::HasValidCompilationAndReplicationAttributes(*op)))
      return WalkResult::interrupt();
    if (op->hasAttr(TF::kReplicationInfoAttr)) {
      op->setAttr(TF::kTpuReplicateAttr, op->getAttr(TF::kReplicationInfoAttr));
      op->removeAttr(TF::kReplicationInfoAttr);
      op->removeAttr(TF::kCompileDeviceTypeAttr);
    }
    return mlir::WalkResult::advance();
  });
  return failure(result.wasInterrupted());

  auto replicate_attr = op.getAttrOfType<StringAttr>(kReplicationInfoAttr);
  auto compile_attr = op.getAttrOfType<StringAttr>(kCompileDeviceTypeAttr);
  if (!replicate_attr && !compile_attr) return success();
  if (!replicate_attr || !compile_attr)
    return op.emitOpError() << "is expected to have either both or none of '"
                            << kReplicationInfoAttr << "' and '"

 public:
  void runOnOperation() override {
    auto traverse_op = [&](Operation* op, tf_device::ClusterOp tpu_cluster) {
      if (isa<tf_device::ClusterOp>(op)) return WalkResult::advance();
      op->removeAttr(TF::kReplicationInfoAttr);
      op->removeAttr(TF::kCompileDeviceTypeAttr);
      // This attribute is used for op colocation. Since all ops are located
      // on a single device cluster, this private attribute is no longer
      // needed.

  auto walk_result = func_op->walk([&](mlir::Operation* op) {
    // Convert `_tpu_replicate`.
    if (op->hasAttr(TF::kTpuReplicateAttr)) {
      op->setAttr(tensorflow::kReplicationInfoAttr,
                  op->getAttr(TF::kTpuReplicateAttr));
      op->removeAttr(tensorflow::kTpuReplicateAttr);
      op->setAttr(tensorflow::kCompileDeviceTypeAttr,

    if (!metadata_op) continue;

    NamedAttrList attrs(metadata_op->getAttrDictionary());

    // Missing or bad `_replication_info` attribute.
    auto replication_info_attr = attrs.get(mlir::TF::kReplicationInfoAttr);
    if (!replication_info_attr)
      return metadata_op.emitError() << kBadReplicateInfoAttrMsg;

    auto replication_info_attr_str =
        mlir::dyn_cast<StringAttr>(replication_info_attr);

    Builder builder(cluster_func_op->getContext());
    // Skip non-tpu device cluster_func.
    auto cluster_id =
        cluster_func_op->getAttrOfType<StringAttr>(TF::kReplicationInfoAttr);
    if (!cluster_id) return WalkResult::advance();

    llvm::SmallVector<int, 4> dynamic_shape_arg_index;

    // Traverse the operands of the cluster func op and find which operand

      }
      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`).
        islands_to_outline.push_back(island_op);
        break;
      }
    }
    return WalkResult::advance();
  });

    "_cluster_outlined_function_name";
// Marks a node for replication. The attribute value indicates the replication
// metadata op.
inline constexpr llvm::StringRef kReplicationInfoAttr = "_replication_info";
// Marks a node for XLA-TPU compilation. The attribute value indicates the
// associated compilation cluster and replication metadata op.

      return true;
    }
  }
  return false;
}

bool IsTPUOp(mlir::Operation* op) {
  return op->hasAttr(TF::kReplicationInfoAttr);
}

StringAttr GetReplicationAttr(mlir::Operation* op) {
  return op->getAttrOfType<StringAttr>(TF::kReplicationInfoAttr);
}

StringAttr GetReplicationAttr(TF::TPUCompilationResultOp op) {
  // Special case for getting the replication region for

    // Op does not belong to a TPU cluster.
    return std::nullopt;
  }
  // Op belongs to a TPU cluster.
  if (auto replication_info =
          op->getAttrOfType<StringAttr>(TF::kReplicationInfoAttr)) {
    // Borrow cluster name from replication info.
    return replication_info.getValue();
  }
  // Use special cluster name for non-replicated case.
  return kNoReplicationCluster;
}