struct CanonicalizeCompileAndReplicateAttributesPass
    : public impl::CanonicalizeCompileAndReplicateAttributesPassBase<
          CanonicalizeCompileAndReplicateAttributesPass> {
  void runOnOperation() override;
};

void CanonicalizeCompileAndReplicateAttributesPass::runOnOperation() {
  func::FuncOp func_op = getOperation();
  ModuleOp module_op = func_op->getParentOfType<ModuleOp>();
  mlir::OpBuilder builder(module_op.getContext());

#include "tensorflow/compiler/mlir/tensorflow/transforms/tf_passes.h.inc"

struct ReplicaIDToDeviceOrdinalPass
    : public impl::ReplicaIDToDeviceOrdinalPassBase<
          ReplicaIDToDeviceOrdinalPass> {
  void runOnOperation() override;
};

// Returns whether op requires `device_ordinal` attribute.
bool RequiresDeviceOrdinalAttribute(Operation* op) {
  return (llvm::isa<TF::EnqueueTPUEmbeddingSparseTensorBatchOp,

#include "tensorflow/compiler/mlir/tensorflow/transforms/tf_device_passes.h.inc"

struct DeviceAttributeToLaunch
    : public impl::DeviceAttributeToLaunchPassBase<DeviceAttributeToLaunch> {
  void runOnOperation() override;
};

void WrapOpInLaunch(Operation* op, llvm::StringRef device) {
  OpBuilder builder(op);

  auto launch_op = builder.create<tf_device::LaunchOp>(
      op->getLoc(), builder.getStringAttr(device),

// XlaInlineDeviceOps Pass will inline cluster op based in the parent region.
struct XlaInlineDeviceOpsPass
    : public impl::XlaInlineDeviceOpsPassBase<XlaInlineDeviceOpsPass> {
  void runOnOperation() override;
};

void InlineDeviceOp(tf_device::ClusterOp cluster_op) {
  auto& block = cluster_op.GetBody();

  llvm::SmallVector<Operation*, 4> non_terminator_ops;

          VerifyInputDialectToExecutorPass> {
 public:
  void runOnOperation() override;
};

bool IsTfDeviceClusterFuncOp(Operation* op) {
  std::string kClusterFuncOpName = "tf_device.cluster_func";
  return op->getName().getStringRef().str() == kClusterFuncOpName;
}

void VerifyInputDialectToExecutorPass::runOnOperation() {
  Operation* func_op = getOperation();

 public:
  explicit TFXLADeviceSpecificTransforms(std::optional<StringRef> device_type) {
    if (device_type.has_value()) {
      device_type_ = device_type.value().str();
    }
  }
  void runOnOperation() override;

 private:
  LogicalResult ConvertGetAlgOp(TF::StatelessRandomGetAlgOp get_alg_op);
};

LogicalResult TFXLADeviceSpecificTransforms::ConvertGetAlgOp(
    TF::StatelessRandomGetAlgOp get_alg_op) {

class VerifyNoOutsideCompilationMarkersPass
    : public impl::VerifyNoOutsideCompilationMarkersPassBase<
          VerifyNoOutsideCompilationMarkersPass> {
 public:
  void runOnOperation() override;
};

bool IsLaunchOp(Operation& op) {
  return dyn_cast<tf_device::LaunchOp>(op) != nullptr;
}

bool IsDeviceClusterOp(Operation& op) {
  return dyn_cast<tf_device::ClusterOp>(op) != nullptr;
}

struct ExecutorDialectToFunctionalConversion
    : public impl::ExecutorDialectToFunctionalPassBase<
          ExecutorDialectToFunctionalConversion> {
  void runOnOperation() override;
};

// Extracts inner ops of tf_executor.island ops in a tf_executor.graph, in the
// order of ops in tf_executor.graph.
LogicalResult LiftIslandOpInnerOpsFromGraph(tf_executor::GraphOp graph) {

      }
    }
  }
  return predecessors;
}

struct TPUColocateSplits
    : public impl::TPUColocateSplitsPassBase<TPUColocateSplits> {
  void runOnOperation() override;
};

void TPUColocateSplits::runOnOperation() {
  getOperation().walk([&](Operation* op) {
    if (auto split = llvm::dyn_cast<TF::SplitOp>(op)) {
      if (HasDevice(split) || split->getAttrOfType<ArrayAttr>(kClassAttr))

struct ConvertToLegacyCompileAndReplicateAttributesPass
    : public impl::ConvertToLegacyCompileAndReplicateAttributesPassBase<
          ConvertToLegacyCompileAndReplicateAttributesPass> {
  void runOnOperation() override;
};

LogicalResult ConvertToLegacyAttributes(func::FuncOp func_op) {
  auto result = func_op->walk([&](mlir::Operation* op) {
    if (failed(TF::HasValidCompilationAndReplicationAttributes(*op)))