1 | #pragma once |
2 | |
3 | #include <torch/csrc/Export.h> |
4 | #include <torch/csrc/jit/api/module.h> |
5 | |
6 | #include <memory> |
7 | #include <string> |
8 | |
9 | namespace torch { |
10 | namespace jit { |
11 | |
12 | /// Compiles script code into an executable graph. |
13 | /// |
14 | /// Takes a string containing functions in script syntax and compiles them into |
15 | /// a module (graph). The returned module provides a `run_method` function |
16 | /// that may be used to invoke the compiled functions. |
17 | /// |
18 | /// For example: |
19 | /// \rst |
20 | /// .. code-block:: cpp |
21 | /// |
22 | /// auto module = torch::jit::compile(R"JIT( |
23 | /// def relu_script(a, b): |
24 | /// return torch.relu(a + b) |
25 | /// def test_while(a, i): |
26 | /// while i < 10: |
27 | /// a += a |
28 | /// i += 1 |
29 | /// return a |
30 | /// )JIT"); |
31 | /// IValue output = module->run_method("relu_script", a, b); |
32 | /// \endrst |
33 | TORCH_API std::shared_ptr<CompilationUnit> compile(const std::string& source); |
34 | |
35 | } // namespace jit |
36 | } // namespace torch |
37 | |