mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
d5df055bbb
Summary: **Summary** This commit adds `torch::jit::RegisterBackend`, an API that allows external backends to be registered for the execution of JIT subgraphs outside the JIT interpreter. In order to register an external backend, one must extend the provided abstract class `PyTorchBackendInterface` and provide two additional functions: one that creates an instance of the aforementioned subclass of `PyTorchBackendInterface`, and another that preprocesses a `ScriptModule` so that it can run on the backend. Then, a `ScriptModule` that can compile and execute a given JIT subgraph using the functions provided at registration time is generated for each registered backend. **Testing** This commit adds a unit test that uses a minimal test backend to make sure that the registration endpoint and generated `ScriptModule` work. ``` $ python test/test_jit.py TestBackends Fail to import hypothesis in common_utils, tests are not derandomized . ---------------------------------------------------------------------- Ran 1 test in 0.183s OK ``` Pull Request resolved: https://github.com/pytorch/pytorch/pull/35833 Differential Revision: D21231955 Pulled By: SplitInfinity fbshipit-source-id: 452db1123d0e5d83f97fe5da8a00fdfdb50dbef9
12 lines
261 B
C++
12 lines
261 B
C++
#include <torch/csrc/jit/backends/backend_init.h>
|
|
#include <torch/csrc/jit/backends/test_backend.h>
|
|
|
|
namespace torch {
|
|
namespace jit {
|
|
|
|
void initJitBackendBindings(PyObject* module) {
|
|
initTestBackendBindings(module);
|
|
}
|
|
} // namespace jit
|
|
} // namespace torch
|