default_device_ = std::string(default_device);
  }

  void runOnOperation() override {
    Builder builder(&getContext());
    Dialect* tf = getContext().getLoadedDialect<TensorFlowDialect>();
    getOperation().walk([&](Operation* op) {
      if (auto device_attr = op->getAttrOfType<StringAttr>(kDeviceAttr)) {
        // We assign default device to ops with device attribute that is empty.

  op->removeAttr(StringAttr::get(context, kDeviceAttr));
  op->moveBefore(return_op);
}

void DeviceAttributeToLaunch::runOnOperation() {
  const Dialect* tf_dialect = getContext().getLoadedDialect("tf");

  getOperation().walk([&](Operation* op) {
    if (op->getDialect() != tf_dialect) return WalkResult::advance();
    if (auto device = op->getAttrOfType<StringAttr>(kDeviceAttr)) {

          op->hasAttr(kDeviceOrdinalAttr) && op->hasAttr(kReplicaIdAttr));
}

void ReplicaIDToDeviceOrdinalPass::runOnOperation() {
  const Dialect* tf_dialect = getContext().getLoadedDialect("tf");
  if (!tf_dialect) {
    getOperation().emitError() << "'tf' dialect is not registered";
    return signalPassFailure();
  }

  // Get the number of devices per host.
  int device_num = 0;

      body.end());

  launch.erase();

  return success();
}

void LaunchToDeviceAttributePass::runOnOperation() {
  const Dialect* tf_dialect = getContext().getLoadedDialect("tf");
  if (!tf_dialect) {
    getOperation().emitError() << "'tf' dialect is not registered";
    return signalPassFailure();
  }

}  // namespace

void populateCanonicalizationPatterns(func::FuncOp func,
                                      RewritePatternSet &patterns) {
  MLIRContext *context = func.getContext();
  mlir::Dialect *tf = context->getLoadedDialect<mlir::TF::TensorFlowDialect>();
  // Load all official canonicalization patterns. Here we skip the
  // canonicalization of the ops in the tf dialect, because they couldn't

  RewritePatternSet patterns(ctx);
  populateWithGenerated(patterns);
  (void)applyPatternsAndFoldGreedily(func, std::move(patterns));

  // Convert unsupported ops to Flex ops.
  auto tf_dialect = ctx->getLoadedDialect<TF::TensorFlowDialect>();
  func.walk([&](Operation *op) {
    if (op->getDialect() != tf_dialect) return;
    if (IsAllowListedOp(op)) return;
    if (op->hasAttr(kNoFallbackAttr)) {

  FunctionLibraryDefinition temp_flib_def(OpRegistry::Global(),
                                          FunctionDefLibrary());
  llvm::SmallDenseSet<FuncOp> visited_functions;
  auto tf_dialect = module.getContext()->getLoadedDialect("tf");
  // Construct SymbolTable to enable cheap function lookups. The cost
  // of constructing the table is offset by the number of queries.
  SymbolTable symbol_table(module);

  FunctionLibraryDefinition temp_flib_def(OpRegistry::Global(),
                                          FunctionDefLibrary());
  llvm::SmallDenseSet<FuncOp> visited_functions;
  auto tf_dialect = module.getContext()->getLoadedDialect("tf");
  // Construct SymbolTable to enable cheap function lookups. The cost
  // of constructing the table is offset by the number of queries.
  SymbolTable symbol_table(module);

      }
    }
  }

  island_op.erase();
  return success();
}

void ReplicateToIslandPass::runOnOperation() {
  const Dialect* tf_dialect = getContext().getLoadedDialect("tf");
  if (!tf_dialect) {
    getOperation().emitError() << "'tf' dialect is not registered";
    return signalPassFailure();
  }

  // Find islands with a single `tf_device.replicate` and create individual

              << " outside compilation regions.";
  }
}

void MarkOpsForOutsideCompilation::runOnOperation() {
  auto module = getOperation();
  const Dialect* tf_dialect = getContext().getLoadedDialect("tf");
  if (!tf_dialect) {
    getOperation().emitError() << "'tf' dialect is not registered";
    return signalPassFailure();
  }
  RewritePatternSet patterns(&getContext());