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"; }