1 #include <torch/csrc/jit/passes/lower_grad_of.h>
2
3 #include <torch/csrc/jit/jit_log.h>
4
5 namespace torch::jit {
6
LowerGradOf(Graph & g)7 void LowerGradOf(Graph& g) {
8 for (auto it = g.nodes().begin(); it != g.nodes().end(); ++it) {
9 if (it->kind() == prim::GradOf) {
10 // if any_defined(inputs):
11 // outputs = <original_computation>
12 // else:
13 // outputs = autograd zero tensors
14 WithInsertPoint guard(*it);
15 auto cond = g.insertNode(g.create(prim::AutogradAnyNonZero, it->inputs()))
16 ->output()
17 ->setType(IntType::get());
18 auto if_stat =
19 g.insertNode(g.create(prim::If, {cond}, it->outputs().size()));
20 if_stat->addBlock()->cloneFrom(
21 it->blocks().at(0), [](Value* v) { return v; });
22 auto else_block = if_stat->addBlock();
23 auto undef = g.createAutogradZero()
24 ->insertBefore(else_block->return_node())
25 ->output();
26 for (size_t i = 0; i < it->outputs().size(); ++i) {
27 // the else block returns a tensor for each of the outputs of the GradOf
28 // i.e. assuming that all the outputs are tensors. This might not be
29 // true, e.g. backward for cat() returns a list of gradient tensors.
30 // This is fixed in DifferentiableGraphBackward, where the list sizes
31 // are stored during the forward pass, and then undefined tensors are
32 // turned into lists of undefined tensors where necessary.
33 else_block->registerOutput(undef);
34 if_stat->outputs().at(i)->copyMetadata(it->outputs().at(i));
35 }
36 GRAPH_UPDATE("Replacing ", getHeader(*it), " with ", getHeader(if_stat));
37 it->replaceAllUsesWith(if_stat);
38 it.destroyCurrent();
39 }
40 }
41 }
42
43 } // namespace torch::jit
44