// Returns true if the op contains a `CalibrationStatisticsSaverOp`.
bool ContainCalibrationStatisticsSaverOp(Operation* op) {
  // Check the region for CaseRegionOp, IfRegionOp and WhileRegionOp.
  for (Region& region : op->getRegions()) {
    if (!region.getOps<TF::CalibrationStatisticsSaverOp>().empty()) {
      return true;
    }
  }

  SymbolTable symbol_table(op->getParentOfType<ModuleOp>());

      op->getAttrs(), op->getPropertiesStorage(), op->getSuccessors(),
      op->getNumRegions());
  builder.insert(new_op);

  // Move region bodies to the new operation.
  for (auto it : llvm::zip(op->getRegions(), new_op->getRegions())) {
    Region &old_region = std::get<0>(it);
    Region &new_region = std::get<1>(it);
    new_region.takeBody(old_region);
  }

  // Replace used results and erase the old op.

  added_call_op->setAttr(kInterfaceNameAttr, interface_name);

  StringAttr device = mlir::cast<StringAttr>(
      added_func_op->getRegion(0).getBlocks().front().front().getAttr(kDevice));
  StringAttr inference_type = mlir::cast<StringAttr>(
      added_func_op->getRegion(0).getBlocks().front().front().getAttr(
          kInferenceType));
  added_call_op->setAttr(kDevice, device);

  llvm::SmallVector<Value, 4> extern_values =
      CollectExternValues(case_region.getRegions());

  // These external values need to be added as inputs to the generated Case. The
  // order is determined by the order of these values the `extern_vales`.

  for (const auto& item : llvm::enumerate(case_region.getRegions())) {
    std::string branch_name =

      op_->getAttrs(), op_->getPropertiesStorage(), op_->getSuccessors(),
      op_->getNumRegions());
  builder.insert(new_op);

  // Move regions to the new op.
  for (auto it : llvm::zip(op_->getRegions(), new_op->getRegions())) {
    Region& old_region = std::get<0>(it);
    Region& new_region = std::get<1>(it);
    new_region.takeBody(old_region);
  }

  // Insert stores to all written resources.

        PropagatePotentiallyWrittenUpFromCallee(func.getRegion(),
                                                call.getArgOperands());
      }
      return;
    }
    if (auto if_op = dyn_cast<TF::IfOp>(op)) {
      for (auto callee : {if_op.then_function(), if_op.else_function()}) {
        PropagatePotentiallyWrittenUpFromCallee(callee.getRegion(),
                                                if_op.getInput());

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() {

                                    cluster_result_types.size() *
                                        num_cores_per_replica);
  for (mlir::Region& region : old_parallel_execute.getRegions()) {
    if (!llvm::isa<tf_device::ClusterFuncOp>(region.front().front())) {
      for (Type t : region.front().front().getResultTypes())
        concatenated_output_types.emplace_back(t);
    }
  }

      }
    } else if (isa<tf_device::LaunchOp, tf_device::ClusterOp,
                   tf_executor::IslandOp, tf_executor::GraphOp, IfRegionOp,
                   CaseRegionOp>(op)) {
      for (Region& region : op->getRegions()) {
        AddRegionSideEffectsForOp(region, op);
      }
    } else if (auto xla_call_module_op = dyn_cast<XlaCallModuleOp>(op)) {
      for (auto func_symbol : xla_call_module_op.getFunctionList().getAsRange<

    // while creating the tf operations.
    new_types.pop_back();

    llvm::SmallVector<std::unique_ptr<Region>, 1> new_regions;
    for (auto &region : op->getRegions()) {
      new_regions.push_back(std::make_unique<Region>());
      new_regions.back()->takeBody(region);
    }

    llvm::SmallVector<NamedAttribute, 4> attrs;