1 #include <c10/util/irange.h>
2 #include <torch/script.h>
3
4 #include "op.h"
5
6 #include <cstddef>
7 #include <string>
8
custom_op(torch::Tensor tensor,double scalar,int64_t repeat)9 torch::List<torch::Tensor> custom_op(
10 torch::Tensor tensor,
11 double scalar,
12 int64_t repeat) {
13 torch::List<torch::Tensor> output;
14 output.reserve(repeat);
15 for (const auto i : c10::irange(repeat)) {
16 (void)i; // Suppress unused variable warning
17 output.push_back(tensor * scalar);
18 }
19 return output;
20 }
21
custom_op2(std::string s1,std::string s2)22 int64_t custom_op2(std::string s1, std::string s2) {
23 return s1.compare(s2);
24 }
25
26 struct CustomOpAutogradFunction : public torch::autograd::Function<CustomOpAutogradFunction> {
forwardCustomOpAutogradFunction27 static torch::Tensor forward(
28 torch::autograd::AutogradContext* ctx,
29 torch::Tensor var1,
30 int64_t mul,
31 torch::Tensor var2,
32 std::optional<torch::Tensor> var3) {
33 ctx->saved_data["mul"] = mul;
34 ctx->saved_data["var3_has_value"] = var3.has_value();
35 ctx->save_for_backward({var1, var2});
36 if (var3) {
37 return var1 + mul * var2 + var1 * var2 + var3.value();
38 }
39 return var1 + mul*var2 + var1*var2;
40 }
41
backwardCustomOpAutogradFunction42 static torch::autograd::variable_list backward(torch::autograd::AutogradContext *ctx, torch::autograd::variable_list grad_output) {
43 int mul = ctx->saved_data["mul"].toInt();
44 bool var3_has_value = ctx->saved_data["var3_has_value"].toBool();
45 auto saved = ctx->get_saved_variables();
46 auto var1 = saved[0];
47 auto var2 = saved[1];
48 auto var3_grad = var3_has_value ? grad_output[0] : torch::Tensor();
49 torch::autograd::variable_list output = {
50 grad_output[0] + grad_output[0] * var2,
51 torch::Tensor(),
52 grad_output[0] * mul + grad_output[0] * var1,
53 var3_grad};
54 return output;
55 }
56 };
57
custom_op_with_autograd(torch::Tensor var1,int64_t mul,torch::Tensor var2,std::optional<torch::Tensor> var3)58 torch::Tensor custom_op_with_autograd(
59 torch::Tensor var1,
60 int64_t mul,
61 torch::Tensor var2,
62 std::optional<torch::Tensor> var3) {
63 return CustomOpAutogradFunction::apply(var1, mul, var2, var3);
64 }
65
custom_nonzero(torch::Tensor x)66 torch::Tensor custom_nonzero(torch::Tensor x) {
67 return x.nonzero();
68 }
69
custom_sin(torch::Tensor x)70 torch::Tensor custom_sin(torch::Tensor x) {
71 return x.sin();
72 }
73
74
TORCH_LIBRARY_FRAGMENT(custom,m)75 TORCH_LIBRARY_FRAGMENT(custom, m) {
76 m.impl_abstract_pystub("my_custom_ops2");
77 m.def("op", custom_op);
78 m.def("op2", custom_op2);
79 m.def("op_with_defaults(Tensor tensor, float scalar = 1, int repeat = 1) -> Tensor[]", custom_op);
80 m.def("op_with_autograd(Tensor var1, int mul, Tensor var2, Tensor? var3=None) -> Tensor", custom_op_with_autograd);
81 m.def("sin(Tensor x) -> Tensor");
82 m.def("cos(Tensor x) -> Tensor");
83 }
84
TORCH_LIBRARY_FRAGMENT(custom,m)85 TORCH_LIBRARY_FRAGMENT(custom, m) {
86 m.impl_abstract_pystub("my_custom_ops");
87 m.def("nonzero(Tensor x) -> Tensor");
88 }
89
TORCH_LIBRARY_FRAGMENT(custom,m)90 TORCH_LIBRARY_FRAGMENT(custom, m) {
91 m.impl_abstract_pystub("nonexistent");
92 m.def("asin(Tensor x) -> Tensor");
93 }
94
TORCH_LIBRARY_FRAGMENT(custom,m)95 TORCH_LIBRARY_FRAGMENT(custom, m) {
96 m.def("tan(Tensor x) -> Tensor");
97 }
98
TORCH_LIBRARY_IMPL(custom,CPU,m)99 TORCH_LIBRARY_IMPL(custom, CPU, m) {
100 m.impl("nonzero", &custom_nonzero);
101 m.impl("sin", &custom_sin);
102 m.impl("asin", &at::asin);
103 }
104