loc, cluster_func.getResultTypes(), operands, cluster_func->getAttrs());
    cluster_func.replaceAllUsesWith(new_cluster_func);
    func::FuncOp func = cluster_func.getFuncOp();
    Block& block = func.front();
    for (Value read_operand : read_operands)
      block.addArgument(read_operand.getType(), loc);

    func.setType(FunctionType::get(&getContext(), block.getArgumentTypes(),

    // Populate set_is_same_data_across_replicas
    // Note: this information is duplicated and can be removed from the proto
    // and here once MLIR bridge phase 2 doesn't fallback to the old bridge.
    auto attr = op.getFuncOp().getArgAttrOfType<mlir::BoolAttr>(
        index, replication_attr_name);
    arg->set_is_same_data_across_replicas(attr != nullptr && attr.getValue());

    // Currently only support first dimension to be bounded dynamic.

    MutableOperandRange getArgOperandsMutable() {
      return getArgsMutable();
    }
    // returns the function that this operation will launch.
    func::FuncOp getFuncOp() {
      return SymbolTable::lookupNearestSymbolFrom<func::FuncOp>(*this, getFuncAttr());
    }
    CallInterfaceCallable getCallableForCallee() {
      return getFuncAttr();
    }

  if (!func_result.wasInterrupted() || !cluster_func.has_value()) {
    return;
  }

  // Get the function on device.
  auto device_func = cluster_func->getFuncOp();
  if (!device_func) return;

  TF::Conv2DOp first_conv;
  // A map maps block argument id to the convolutions consumes them.
  llvm::SmallDenseMap<unsigned, std::vector<Conv2DWithBlockSize>>