1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
|
#pragma once
#include <c10/util/irange.h>
#include <torch/csrc/Export.h>
#include <torch/csrc/autograd/function.h>
#include <torch/csrc/autograd/variable.h>
#include <memory>
#include <string>
#include <vector>
namespace torch::autograd {
struct TORCH_API Error : public Node {
Error(std::string msg, edge_list&& next_edges)
: Node(std::move(next_edges)), msg(std::move(msg)) {}
Error(std::string msg) : msg(std::move(msg)) {}
variable_list apply(variable_list&& inputs) override;
void compiled_args(CompiledNodeArgs& args) override;
variable_list apply_with_saved(
const variable_list& inputs,
SwapSavedVariables& saved) override;
std::string msg;
};
// We print grad_fn names in tensor printing. For functions with backward
// NYI, grad_fn=<Error> will be printed if we use Error, which is confusing. So
// special case with a new NotImplemented function here.
struct TORCH_API NotImplemented : public Error {
NotImplemented(const std::string& forward_fn, edge_list&& next_edges)
: Error(
"derivative for " + forward_fn + " is not implemented",
std::move(next_edges)) {}
NotImplemented(const std::string& forward_fn)
: Error("derivative for " + forward_fn + " is not implemented") {}
};
// Identity in forward, Error in backward. Used to implement
// @once_differentiable
struct TORCH_API DelayedError : public Node {
DelayedError(std::string msg, int64_t num_inputs) : msg(std::move(msg)) {
for ([[maybe_unused]] const auto _ [[maybe_unused]] :
c10::irange(num_inputs)) {
add_input_metadata(Node::undefined_input());
}
}
variable_list apply(variable_list&& inputs) override;
std::string msg;
};
struct TORCH_API UndefinedGrad : public Node {
UndefinedGrad() {
add_input_metadata(Node::undefined_input());
}
variable_list apply(variable_list&& inputs) override;
};
struct TORCH_API UndefinedGradBackward : public Node {
UndefinedGradBackward(edge_list&& next_edges) : Node(std::move(next_edges)) {}
UndefinedGradBackward() = default;
variable_list apply(variable_list&& inputs) override;
void compiled_args(CompiledNodeArgs& args) override {}
variable_list apply_with_saved(
const variable_list& inputs,
SwapSavedVariables& saved) override {
return apply(variable_list(inputs));
}
};
struct TORCH_API GraphRoot : public Node {
GraphRoot(edge_list functions, variable_list inputs)
: Node(std::move(functions)), outputs(std::move(inputs)) {
// Ensures calls to stream() on a GraphRoot instance reflect current
// stream(s) on devices of root grad tensors at the time the instance is
// constructed.
for (const auto& t : outputs) {
add_input_metadata(t);
}
}
variable_list apply(variable_list&& inputs) override {
return outputs;
}
void compiled_args(CompiledNodeArgs& args) override;
variable_list apply_with_saved(
const variable_list& inputs,
SwapSavedVariables& saved) override;
variable_list outputs;
};
struct TORCH_API Identity : public Node {
variable_list apply(variable_list&& inputs) override;
};
} // namespace torch::autograd
|
