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
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
|
#pragma once
#include <ATen/core/function.h>
#include <ATen/core/ivalue.h>
#include <ATen/core/stack.h>
#include <torch/csrc/api/include/torch/imethod.h>
#include <torch/csrc/jit/api/function_impl.h>
namespace torch::jit {
using ObjectPtr = c10::intrusive_ptr<c10::ivalue::Object>;
// A method in a module, e.g. f in:
//
// class M(ScriptModule):
// @script_method
// def f(self, x):
// ...
// Note: because Method/Module are exposed to python these
// classes use python method naming conventions
struct TORCH_API Method : public torch::IMethod {
Method(ObjectPtr owner, Function* function);
// the module that contains this method.
Module owner() const;
// the raw objectptr that owns this method, for when the method is owned by a
// torchbind object.
ObjectPtr raw_owner() const;
void run(Stack& stack);
void run(Stack&& stack) {
run(stack);
}
c10::IValue operator()(
std::vector<c10::IValue> stack,
const Kwargs& kwargs = Kwargs()) const override;
// Run method async. Invocation on this function would invokes a JIT
// interpreter that executes ops inline, one by one, on caller's thread. A
// model can utilize async op, i.e. `fork`, to launch an asynchronous task
// which will be launched on provided `taskLauncher`.
c10::intrusive_ptr<c10::ivalue::Future> run_async(
std::vector<c10::IValue> stack,
const Kwargs& kwargs = Kwargs(),
TaskLauncher taskLauncher = at::launch);
std::shared_ptr<Graph> graph() const {
return toGraphFunction(*function_).graph();
}
const std::string& name() const override {
return function_->name();
}
size_t num_inputs() const {
return function_->num_inputs();
}
GraphExecutor& get_executor() {
return toGraphFunction(*function_).get_executor();
}
Function& function() const {
return *function_;
}
private:
void setArgumentNames(std::vector<std::string>&) const override;
// Methods are uniqued onwed by a single module. This raw pointer allows
// looking up the module.
ObjectPtr owner_;
// Underlying unbound function
Function* function_;
};
namespace script {
// We once had a `script::` namespace that was deleted. This is for backcompat
// of the public API; new code should not use this type alias.
using Method = ::torch::jit::Method;
} // namespace script
} // namespace torch::jit
|
