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

      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;
}

    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,