return llvm::cast<func::CallOp>(output_op);
  }
  return nullptr;
}

class PickSubgraphsPass
    : public mlir::PassWrapper<PickSubgraphsPass,
                               mlir::OperationPass<ModuleOp>> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PickSubgraphsPass)

 private:
  llvm::StringRef getArgument() const final { return "tfl-pick-subgraphs"; }

    llvm::ArrayRef<std::string> device_specs);

// Create an instance of ComputeCostPass.
std::unique_ptr<OperationPass<ModuleOp>> CreateComputeCostPass();

// Create an instance of PickSubgraphsPass.
std::unique_ptr<OperationPass<ModuleOp>> CreatePickSubgraphsPass();

// Create an instance of DeviceTransformGPUPass.
std::unique_ptr<OperationPass<func::FuncOp>> CreateDeviceTransformGPUPass();