Caffe2 - C++ API
A deep learning, cross platform ML framework
Related Pages
Modules
Data Structures
Files
C++ API
Python API
GitHub
File List
Globals
test
cpp_extensions
jit_extension.cpp
1
#include <torch/extension.h>
2
3
#include "doubler.h"
4
5
using namespace
at
;
6
7
Tensor
exp_add(
Tensor
x,
Tensor
y);
8
9
Tensor
tanh_add(
Tensor
x,
Tensor
y) {
10
return
x.tanh() + y.tanh();
11
}
12
13
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
14
m.def(
"tanh_add"
, &tanh_add,
"tanh(x) + tanh(y)"
);
15
m.def(
"exp_add"
, &exp_add,
"exp(x) + exp(y)"
);
16
py::class_<Doubler>(m,
"Doubler"
)
17
.def(py::init<int, int>())
18
.def(
"forward"
, &Doubler::forward)
19
.def(
"get"
, &Doubler::get);
20
}
at::Tensor
Definition:
Tensor.h:48
at
Flush-To-Zero and Denormals-Are-Zero mode.
Definition:
AccumulateType.h:17
Generated on Thu Mar 21 2019 13:06:22 for Caffe2 - C++ API by
1.8.11