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
|
#pragma once
#include <pybind11/pybind11.h>
#include "c10/util/Registry.h"
namespace caffe2 {
namespace python {
namespace py = pybind11;
struct PybindAddition {
PybindAddition() {}
PybindAddition(py::module&) {}
virtual ~PybindAddition(){};
};
C10_DECLARE_REGISTRY(PybindAdditionRegistry, PybindAddition, py::module&);
#define REGISTER_PYBIND_ADDITION(funcname) \
namespace { \
struct funcname##Impl : public PybindAddition { \
funcname##Impl(py::module& m) { \
funcname(m); \
} \
}; \
C10_REGISTER_CLASS(PybindAdditionRegistry, funcname##Impl, funcname##Impl); \
}
} // namespace python
} // namespace caffe2
|
