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
|
#include <c10/util/irange.h>
#include <torch/script.h>
#include "op.h"
#include <cstddef>
#include <string>
torch::List<torch::Tensor> custom_op(
torch::Tensor tensor,
double scalar,
int64_t repeat) {
torch::List<torch::Tensor> output;
output.reserve(repeat);
for (const auto i : c10::irange(repeat)) {
(void)i; // Suppress unused variable warning
output.push_back(tensor * scalar);
}
return output;
}
int64_t custom_op2(std::string s1, std::string s2) {
return s1.compare(s2);
}
struct CustomOpAutogradFunction : public torch::autograd::Function<CustomOpAutogradFunction> {
static torch::Tensor forward(
torch::autograd::AutogradContext* ctx,
torch::Tensor var1,
int64_t mul,
torch::Tensor var2,
c10::optional<torch::Tensor> var3) {
ctx->saved_data["mul"] = mul;
ctx->saved_data["var3_has_value"] = var3.has_value();
ctx->save_for_backward({var1, var2});
if (var3) {
return var1 + mul * var2 + var1 * var2 + var3.value();
}
return var1 + mul*var2 + var1*var2;
}
static torch::autograd::variable_list backward(torch::autograd::AutogradContext *ctx, torch::autograd::variable_list grad_output) {
int mul = ctx->saved_data["mul"].toInt();
bool var3_has_value = ctx->saved_data["var3_has_value"].toBool();
auto saved = ctx->get_saved_variables();
auto var1 = saved[0];
auto var2 = saved[1];
auto var3_grad = var3_has_value ? grad_output[0] : torch::Tensor();
torch::autograd::variable_list output = {
grad_output[0] + grad_output[0] * var2,
torch::Tensor(),
grad_output[0] * mul + grad_output[0] * var1,
var3_grad};
return output;
}
};
torch::Tensor custom_op_with_autograd(
torch::Tensor var1,
int64_t mul,
torch::Tensor var2,
c10::optional<torch::Tensor> var3) {
return CustomOpAutogradFunction::apply(var1, mul, var2, var3);
}
TORCH_LIBRARY_FRAGMENT(custom, m) {
m.def("op", custom_op);
m.def("op2", custom_op2);
m.def("op_with_defaults(Tensor tensor, float scalar = 1, int repeat = 1) -> Tensor[]", custom_op);
m.def("op_with_autograd(Tensor var1, int mul, Tensor var2, Tensor? var3=None) -> Tensor", custom_op_with_autograd);
}
|
