}

// Replaces the replication region attribute if it already exists.
void UpdateReplicationAttr(Operation* op, StringAttr attr) {
  if (op->hasAttr(TF::kReplicationInfoAttr)) {
    op->setAttr(TF::kReplicationInfoAttr, attr);
  }
}

// Replaces the replication region attribute if it already exists.
void UpdateReplicationAttr(TF::TPUCompilationResultOp& op, StringAttr attr) {

    if (failed(TF::HasValidCompilationAndReplicationAttributes(*op)))
      return WalkResult::interrupt();
    // Skip non-tpu device cluster_func.
    auto cluster_id = op->getAttrOfType<StringAttr>(TF::kReplicationInfoAttr);
    if (!cluster_id) return WalkResult::advance();

    if (failed(Rewrite(module_name, op, devices.device_names(),
                       compilation_results[cluster_id], &builder,