.getOperation();
           });

  // mlir::TF::AddOp
  py::class_<mlir::TF::AddV2Op>(m, "Tf_AddV2Op")
      .def("create",
           [](mlir::OpBuilder& opb, mlir::Location loc, mlir::Value x,
              mlir::Value y) -> mlir::Operation* {
             return opb.create<mlir::TF::AddV2Op>(loc, x, y).getOperation();
           });

               << "Found call to inline: " << *call_op.getOperation() << "\n");

    auto call_interface = cast<CallOpInterface>(call_op.getOperation());
    auto called_func =
        dyn_cast_or_null<func::FuncOp>(call_interface.resolveCallable());

    if (failed(inlineCall(inliner, call_interface,
                          cast<CallableOpInterface>(called_func.getOperation()),
                          called_func.getCallableRegion(),

        return WalkResult::interrupt();
      host_device = host_device_value;
    }
    cluster_to_host[tpu_cluster.getOperation()] = host_device;
    return tpu_cluster.walk([&](Operation* op) {
      insert_pending_op(op, tpu_cluster);
      return callback(op, tpu_cluster,
                      cluster_to_host[tpu_cluster.getOperation()]);
    });
  });
  if (result.wasInterrupted()) return failure();

      input_shapes_ = input_shapes_vec;
    }

    auto failure_or_converged = InferModuleShape(
        getOperation(), max_iterations_, /*ops_to_skip=*/{}, input_shapes_);
    if (failed(failure_or_converged)) return signalPassFailure();
    if (!failure_or_converged.value()) {
      getOperation().emitError()
          << "shape inference pass did not reach convergence after "
          << max_iterations_;

  if (op->hasAttr("_wraps_compilation")) return 1;
  // Followed by everything else.
  return 0;
}

void MoveTpuCompileToFrontPass::runOnOperation() {
  MarkCompilationOps(getOperation());
  getOperation().walk([](Operation* op) {
    for (Region& region : op->getRegions()) {
      for (Block& block : region.getBlocks()) {
        if (block.empty()) continue;

  return count;
}

void NameAnonymousIteratorsPass::runOnOperation() {
  int count = 1;
  getOperation().walk(
      [&](TF::AnonymousIteratorOp op) { count = replace(op, count); });
  getOperation().walk(
      [&](TF::AnonymousIteratorV2Op op) { count = replace(op, count); });
  getOperation().walk(
      [&](TF::AnonymousIteratorV3Op op) { count = replace(op, count); });
}

}  // namespace

  // SymbolRefAttrs of all ops are present.
  SymbolTable symbol_table = SymbolTable(getOperation());
  llvm::SetVector<func::FuncOp> reachable_funcs;
  for (auto func : getOperation().getOps<func::FuncOp>()) {
    auto walk_result = func.walk([&](func::CallOp op) -> WalkResult {
      if (!symbol_table.lookup<func::FuncOp>(op.getCallee()))
        return getOperation().emitError()
               << func.getName() << " is not in the funcs allowlist";

// order of ops in tf_executor.graph.
LogicalResult LiftIslandOpInnerOpsFromGraph(tf_executor::GraphOp graph) {
  auto graph_position = graph.getOperation()->getIterator();
  Block* parent_block = graph.getOperation()->getBlock();
  for (Operation& op : graph.GetBody().without_terminator()) {
    auto island_op = llvm::dyn_cast<tf_executor::IslandOp>(op);
    if (!island_op)
      return op.emitOpError()

    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.
        if (device_attr.getValue().empty()) {

  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;
  mlir::TF::RuntimeDevices devices;
  if (failed(tensorflow::GetDevicesFromOp(
          getOperation()->getParentOfType<ModuleOp>(), &devices)))
    return signalPassFailure();