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
|
#pragma once
#include <torch/csrc/autograd/function_hook.h>
#include <torch/csrc/dynamo/compiled_autograd.h>
namespace torch::autograd::utils {
// Turns lambda into a torch::autograd::FunctionPostHook.
class LambdaPostHook : public torch::autograd::FunctionPostHook {
using variable_list = std::vector<torch::autograd::Variable>;
using fn_type =
std::function<variable_list(const variable_list&, const variable_list&)>;
using compiled_fn_type = std::function<void(CompiledNodeArgs&)>;
public:
// The lambda function takes as arguments the outputs and inputs of the
// autograd function and can modify the outputs of the autograd function by
// returning a new output if needed.
/* implicit */ LambdaPostHook(fn_type fn) : fn_(std::move(fn)) {}
LambdaPostHook(fn_type fn, compiled_fn_type compiled_fn)
: fn_(std::move(fn)), compiled_fn_(std::move(compiled_fn)) {}
variable_list operator()(
const variable_list& outputs,
const variable_list& inputs) override {
return fn_(outputs, inputs);
}
void compiled_args(CompiledNodeArgs& args) override {}
protected:
std::function<variable_list(const variable_list&, const variable_list&)> fn_;
compiled_fn_type compiled_fn_{};
};
} // namespace torch::autograd::utils
|
