xref: /aosp_15_r20/external/pytorch/torch/csrc/api/include/torch/jit.h (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
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