1#pragma once
2
3#include <torch/csrc/Export.h>
4#include <torch/csrc/jit/ir/ir.h>
5
6#include <cstddef>
7
8namespace torch {
9namespace jit {
10
11// insert GraphExecutor nodes that group together
12// subgraphs that are differentiable by the jit's autodiff passes
13// threshold - minimum number of nodes that will appear in a block
14// returns all differentiable blocks that have been found
15TORCH_API std::vector<Node*> CreateAutodiffSubgraphs(
16 const std::shared_ptr<Graph>& graph,
17 size_t threshold = 2);
18} // namespace jit
19} // namespace torch
20