if (absl::StartsWith(op_name, "tf.") || absl::StartsWith(op_name, "tfl.")) {
      loc_to_name_.emplace(op->getLoc(), op_name);
    }
  };

  for (auto &region : module->getRegions()) {
    region.walk(collectOps);
  }

  pass_name_ = extract_pass_name(pass->getName().str());
  error_collector_->Clear();
}

void ErrorCollectorInstrumentation::runAfterPass(Pass *pass,

void RewriteTPUEmbeddingOps::runOnOperation() {
  func::FuncOp func = getOperation();
  if (failed(RunOnRegion(&func.getBody()))) return signalPassFailure();

  func.walk([&](Operation* op) {
    for (Region& region : op->getRegions()) {
      if (failed(RunOnRegion(&region))) return signalPassFailure();
    }
  });
}

}  // anonymous namespace

std::unique_ptr<OperationPass<func::FuncOp>>
CreateRewriteTPUEmbeddingOpsPass() {

  // _TPUDeviceOrdinalPlaceholder implicitly depends on the replica.
  if (llvm::isa<TF::_TPUDeviceOrdinalPlaceholderOp>(op)) return false;

  bool has_replicate_operands = false;
  visitUsedValuesDefinedAbove(op->getRegions(), [&](OpOperand* operand) {
    if (!ancestor_of_replicate(operand->get().getParentRegion()))
      has_replicate_operands = true;
  });

  return !has_replicate_operands;
}

  for (int i = 0; i < op->getNumRegions(); ++i) {
    state.addRegion();
  }
  Operation* new_op = builder.create(state);
  for (const auto& indexed_regions : llvm::enumerate(op->getRegions())) {
    Region& region = op->getRegion(indexed_regions.index());
    IRMapping mapping;
    indexed_regions.value().cloneInto(&region, mapping);
  }
  int new_position = 0;
  for (auto result : op->getResults()) {

    OperationState state(op->getLoc(), op->getName().getStringRef(), operands,
                         new_results, op->getAttrs(), op->getSuccessors());
    for (Region& region : op->getRegions()) {
      Region& new_region = *state.addRegion();
      rewriter.inlineRegionBefore(region, new_region, new_region.begin());
      if (failed(rewriter.convertRegionTypes(&new_region, *getTypeConverter())))

      continue;
    }
    // Gets the live in values of the `op`
    llvm::SetVector<Value> live_ins(op.operand_begin(), op.operand_end());
    getUsedValuesDefinedAbove(op.getRegions(), live_ins);
    // Inserts if any of the `live_ins` depends on the ops in the cluster.
    if (llvm::any_of(live_ins, [&](Value value) {
          Operation* defining_op = value.getDefiningOp();
          if (!defining_op) {

        TF::TensorFlowDialect::getDialectNamespace())
      return failure();

    // Refining types may have implications to the attached regions or symbol
    // references so do not update such ops.
    if (!op->getRegions().empty() || HasSymbolRefAttr(op)) return failure();

    IRMapping mapper;
    bool has_type_change = false;
    for (auto [original, updated] : llvm::zip(op->getOperands(), operands)) {

        new URI("s3://somebucket.au.s3.eu-central-1.amazonaws.com/a/b/file.txt")      | Optional.of(getRegion(Regions.EU_CENTRAL_1))                        | 'somebucket.au' | 'a/b/file.txt'

      replicate_op->getLoc(), new_result_types, replicate_op->getOperands(),
      replicate_op->getAttrs());

  // Move region to the new op.
  new_replicate_op.getRegion().takeBody(replicate_op.getRegion());

  // Replace all old uses with new op results.
  int old_num_results = replicate_op->getNumResults();
  replicate_op->replaceAllUsesWith(

            String key = matcher.group(4);
            Region derivedRegion;
            if (region.equals("external-1")) {
                derivedRegion = Region.getRegion(Regions.US_EAST_1);
            } else {
                derivedRegion = RegionUtils.getRegion(region);
            }

            this.region = Optional.of(derivedRegion);
            this.bucketName = bucketName;
            this.key = key;