Files
pytorch/torch/csrc/jit/python/script_init.cpp
Pavithran Ramachandran fc2cf3d26f Back out "Revert D34805092: Extend _save_for_mobile and _load_for_mobile to support flatbuffer format; Default format is pickle + Change buck targets to support only pickle and pickle + flatbuffer for migration" (#74594)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/74594

Extending `_save_for_mobile` and `_load_for_mobile` to support faltbuffer format with additional optional argument which is set to pick pickle by default.

Adding new binary target with suffix `_pickle_and_flatbuffer` to help migration.

Size test in D34909502 shows the size has regressed by ~40K but after removing pickle and comparing lite_predictors we have ~120K size measure that we will achieve when deprecating pickle and moving to flatbuffer

**BEFORE:**

```lang=mermaid
graph TD;
    torch_core-->torch_mobile_deserialize;

    torch_mobile_core-->torch_mobile_deserialize;

    jit_module_saving-->torch_core;
    jit_module_saving-->torch_mobile_core;

    torch_mobile_deserialize-->caffe2_serialize;
    torch_mobile_deserialize-->torch_mobile_module;

    caffe2_serialize-->miniz;

    flatbuffer_loader-->mobile_bytecode;
    flatbuffer_serializer-->mobile_bytecode;

    mobile_bytecode-->flatbuffer_2.0;

    flatbuffer_loader-->torch_mobile_module;
    flatbuffer_serializer-->torch_mobile_module;
```

**AFTER:**
```lang=mermaid
graph TD;
    torch_core-->torch_mobile_deserialize;

    torch_mobile_core-->torch_mobile_deserialize;

    jit_module_saving-->torch_core;
    jit_module_saving-->torch_mobile_core;

    torch_mobile_deserialize-->caffe2_serialize;
    torch_mobile_deserialize-->torch_mobile_module;

    caffe2_serialize-->miniz;

    flatbuffer_loader-->mobile_bytecode;
    flatbuffer_serializer-->mobile_bytecode;

    mobile_bytecode-->flatbuffer_2.0;

    torch_mobile_deserialize_pickle_and_flatbuffer-->|new| flatbuffer_loader;
    torch_mobile_deserialize_pickle_and_flatbuffer-->|new| torch_mobile_deserialize;
    torch_mobile_core_pickle_and_flatbuffer-->|new| torch_mobile_deserialize_pickle_and_flatbuffer;
    torch_core_pickle_and_flatbuffer-->|new| torch_mobile_deserialize_pickle_and_flatbuffer;

    jit_module_saving_pickle_and_flatbuffer-->|new| torch_core_pickle_and_flatbuffer;
    jit_module_saving_pickle_and_flatbuffer-->|new| torch_mobile_core_pickle_and_flatbuffer;

    flatbuffer_serializer-->torch_mobile_module;

    jit_module_saving_pickle_and_flatbuffer-->|new|jit_module_saving;
    jit_module_saving_pickle_and_flatbuffer-->|new|flatbuffer_serializer;

    flatbuffer_loader-->torch_mobile_module;
```

Original commit changeset: 780dfb6fd6ba

Original Phabricator Diff: D34805092 (284b2b7135)
ghstack-source-id: 152044801

(Note: this ignores all push blocking failures!)

Test Plan:
CI

```
~/fbsource/fbcode] cd ~/fbsource/fbcode/ && buck test -c fbcode.caffe2_enable_flatbuffer=1 //caffe2/test/cpp/jit:jit  -- FlatbufferTest.ExtraFiles
Parsing buck files: finished in 0.9 sec
Building: finished in 5.3 sec (100%) 12992/54304 jobs, 0/54304 updated
  Total time: 6.2 sec
More details at https://www.internalfb.com/intern/buck/build/2b387fff-f813-4cfa-b53f-eb2378630d4e
BUILD SUCCEEDED
Tpx test run coordinator for Facebook. See https://fburl.com/tpx for details.
Running with tpx session id: f93a84d6-e7ce-41a0-a97f-0ef3fa6d199d
Trace available for this run at /tmp/tpx-20220323-134108.766518-f93a84d6-e7ce-41a0-a97f-0ef3fa6d199d/trace.log
RemoteExecution session id: reSessionID-f93a84d6-e7ce-41a0-a97f-0ef3fa6d199d-tpx
Started reporting to test run: https://www.internalfb.com/intern/testinfra/testrun/4503599723101693
    ✓ ListingSuccess: caffe2/test/cpp/jit:jit : 486 tests discovered (19.122)
    ✓ Pass: caffe2/test/cpp/jit:jit - FlatbufferTest.ExtraFiles (0.187)
Summary
  Pass: 1
  ListingSuccess: 1
If you need help understanding your runs, please follow the wiki: https://fburl.com/posting_in_tpx_users
Finished test run: https://www.internalfb.com/intern/testinfra/testrun/4503599723101693
```

Similar Build Deps Dags

```
[pavithran@devvm5216.vll0 /data/users/pavithran/fbsource] buck query 'allpaths(//xplat/caffe2:torch_mobile_all_ops_pickle_and_flatbuffer, //xplat/caffe2:torch_mobile_deserialize_pickle_and_flatbuffer)' --output-format dot-compact  | pastry
P486770901: https://www.internalfb.com/intern/paste/P486770901/

[pavithran@devvm5216.vll0 /data/users/pavithran/fbsource] buck query 'allpaths(//xplat/caffe2:torch_mobile_all_ops, //xplat/caffe2:torch_mobile_deserialize)' --output-format dot-compact  | pastry
P486771278: https://www.internalfb.com/intern/paste/P486771278/
```

pickle_and_flatbuffer: https://www.internalfb.com/intern/dgw/graph/?build_id=P486770901
pickle: https://www.internalfb.com/intern/dgw/graph/?build_id=P486771278

Reviewed By: iseeyuan

Differential Revision: D35067157

fbshipit-source-id: 9044259c17a2e0da79bd6aedb28efbdfd57e23e0
(cherry picked from commit f738069ec3a72e79da56172741d027de514e9e5f)
2022-03-24 21:51:05 +00:00

2235 lines
81 KiB
C++

#include <pybind11/detail/common.h>
#include <pybind11/pytypes.h>
#include <torch/csrc/jit/api/object.h>
#include <torch/csrc/jit/python/script_init.h>
#include <caffe2/serialize/versions.h>
#include <torch/csrc/Device.h>
#include <torch/csrc/DynamicTypes.h>
#include <torch/csrc/jit/api/module.h>
#include <torch/csrc/jit/frontend/ir_emitter.h>
#include <torch/csrc/jit/frontend/sugared_value.h>
#include <torch/csrc/jit/mobile/code.h>
#include <torch/csrc/jit/mobile/compatibility/backport.h>
#include <torch/csrc/jit/mobile/compatibility/model_compatibility.h>
#include <torch/csrc/jit/mobile/import.h>
#include <torch/csrc/jit/mobile/module.h>
#include <torch/csrc/jit/operator_upgraders/upgraders.h>
#include <torch/csrc/jit/operator_upgraders/upgraders_entry.h>
#include <torch/csrc/jit/operator_upgraders/upgraders_guard.h>
#include <torch/csrc/jit/operator_upgraders/utils.h>
#include <torch/csrc/jit/operator_upgraders/version_map.h>
#include <torch/csrc/jit/python/module_python.h>
#include <torch/csrc/jit/python/python_ivalue.h>
#include <torch/csrc/jit/python/python_sugared_value.h>
#include <torch/csrc/jit/serialization/export_bytecode.h>
#include <torch/csrc/jit/serialization/import.h>
#include <torch/csrc/jit/testing/file_check.h>
#include <c10/util/intrusive_ptr.h>
#include <c10/util/irange.h>
#include <torch/csrc/jit/frontend/parser.h>
#include <torch/csrc/jit/frontend/tracer.h>
#include <torch/csrc/jit/ir/constants.h>
#include <torch/csrc/jit/ir/irparser.h>
#include <torch/csrc/jit/passes/inliner.h>
#include <torch/csrc/jit/passes/shape_analysis.h>
#include <torch/csrc/jit/python/pybind_utils.h>
#include <torch/csrc/jit/python/python_dict.h>
#include <torch/csrc/jit/python/python_list.h>
#include <torch/csrc/jit/python/python_tracer.h>
#include <torch/csrc/jit/runtime/graph_executor.h>
#include <torch/csrc/jit/runtime/instruction.h>
#include <torch/csrc/jit/runtime/interpreter.h>
#include <torch/csrc/jit/runtime/logging.h>
#include <torch/csrc/jit/serialization/import_source.h>
#include <torch/csrc/jit/serialization/python_print.h>
#include <torch/csrc/jit/testing/hooks_for_testing.h>
#include <torch/csrc/api/include/torch/ordered_dict.h>
#include <ATen/ATen.h>
#include <ATen/core/function_schema.h>
#include <ATen/core/ivalue.h>
#include <ATen/core/qualified_name.h>
#include <pybind11/functional.h>
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
#include <pybind11/stl_bind.h>
#include <chrono>
#include <cstddef>
#include <memory>
#include <sstream>
#include <string>
#include <tuple>
#include <utility>
#include <vector>
namespace torch {
namespace jit {
using ::c10::Argument;
using ::c10::FunctionSchema;
using ResolutionCallback = std::function<py::object(std::string)>;
using FunctionDefaults = std::unordered_map<std::string, py::object>;
using ClassMethodDefaults = std::unordered_map<std::string, FunctionDefaults>;
namespace {
// A resolver that will inspect the outer Python scope to find `name`.
struct PythonResolver : public Resolver {
explicit PythonResolver(ResolutionCallback rcb) : rcb_(std::move(rcb)) {}
/**
* While compiling classes, the class type we're compiling will not be
* available in Python, since we haven't fowner_ defining the class yet. So
* in order to make the class type available to its own methods, we need to
* explicitly resolve it.
*
* @param rcb Python function to resolve a name to its Python object in the
* enclosing scope
* @param classname The unqualified classname of the class currently being
* compiled.
* @param classType The class's type.
*/
explicit PythonResolver(
ResolutionCallback rcb,
std::string classname,
ClassTypePtr classType)
: rcb_(std::move(rcb)),
classname_(std::move(classname)),
classType_(std::move(classType)) {}
std::shared_ptr<SugaredValue> resolveValue(
const std::string& name,
GraphFunction& m,
const SourceRange& loc) override {
pybind11::gil_scoped_acquire ag;
py::object obj = rcb_(name);
if (obj.is(py::none())) {
return nullptr;
}
return toSugaredValue(obj, m, loc);
}
static bool isNamedTupleClass(py::object obj) {
auto tuple_type = reinterpret_cast<PyObject*>(&PyTuple_Type);
return PyObject_IsSubclass(obj.ptr(), tuple_type) &&
py::hasattr(obj, "_fields");
}
TypePtr resolveTypeFromObject(const py::object& obj, const SourceRange& loc) {
if (py::isinstance<ScriptClass>(obj)) {
auto script_class = py::cast<ScriptClass>(obj);
return script_class.class_type_.type_;
}
py::bool_ isClass = py::module::import("inspect").attr("isclass")(obj);
if (!py::cast<bool>(isClass)) {
return nullptr;
}
if (isNamedTupleClass(obj)) {
return registerNamedTuple(obj, loc);
}
auto qualifiedName = c10::QualifiedName(
py::cast<std::string>(py::module::import("torch._jit_internal")
.attr("_qualified_name")(obj)));
return get_python_cu()->get_type(qualifiedName);
}
TypePtr resolveType(const std::string& name, const SourceRange& loc)
override {
if (classType_ && name == classname_) {
return classType_;
}
pybind11::gil_scoped_acquire ag;
py::object obj = rcb_(name);
if (obj.is(py::none())) {
return nullptr;
}
auto annotation_type = py::module::import("torch.jit.annotations")
.attr("try_ann_to_type")(obj, loc);
if (!annotation_type.is_none()) {
return py::cast<TypePtr>(annotation_type);
}
return resolveTypeFromObject(obj, loc);
}
private:
ResolutionCallback rcb_;
std::string classname_;
ClassTypePtr classType_;
};
std::shared_ptr<PythonResolver> pythonResolver(const ResolutionCallback& rcb) {
return std::make_shared<PythonResolver>(rcb);
}
std::shared_ptr<PythonResolver> pythonResolver(
const ResolutionCallback& rcb,
std::string classname,
ClassTypePtr classType) {
return std::make_shared<PythonResolver>(
rcb, std::move(classname), std::move(classType));
}
void checkOverloadDecl(const Decl& new_decl, const Decl& old_decl) {
const auto& new_params = new_decl.params();
const auto& old_params = old_decl.params();
// TODO. same number of parameters not strictly necessary.
TORCH_INTERNAL_ASSERT(
new_params.size() == old_params.size(),
"Overload must have same number of parameters\n",
new_decl.range(),
old_decl.range());
for (const auto i : c10::irange(new_decl.params().size())) {
TORCH_INTERNAL_ASSERT(
new_params[i].ident().name() == old_params[i].ident().name(),
"Overload parameters must have the same names\n",
new_params[i].ident(),
old_params[i].ident());
}
}
c10::optional<IValue> tryCalculateDefaultParam(
const Argument& arg,
const py::object& def_value) {
auto n = arg.N();
auto list_type = arg.type()->cast<ListType>();
try {
if (n && *n > 0 && list_type) {
// BroadcastingList, allow default values T for arg types List[T]
return toIValue(def_value, list_type->getElementType());
} else {
return toIValue(def_value, arg.type());
}
} catch (...) {
return c10::nullopt;
}
}
// An overloaded function may have a default that does not subtype all overloads
// @overload
// def foo(x: str)
// def foo(x=1)
FunctionDefaults calcOverloadedFunctionDefaults(
const FunctionSchema& schema,
const FunctionDefaults& defaults) {
FunctionDefaults updated_defaults;
for (const auto& arg : schema.arguments()) {
const std::string& arg_name = arg.name();
auto value = defaults.find(arg_name);
if (value == defaults.end()) {
continue;
}
auto maybe_ivalue = tryCalculateDefaultParam(arg, value->second);
if (maybe_ivalue) {
updated_defaults[arg_name] = value->second;
}
}
return updated_defaults;
}
} // namespace
bool checkMutableFunctionDefault(const py::object& def_arg) {
if (py::isinstance<py::list>(def_arg) || py::isinstance<py::dict>(def_arg)) {
return true;
}
if (py::isinstance<py::tuple>(def_arg)) {
auto pytuple = def_arg.cast<py::tuple>();
for (py::handle t : pytuple) {
py::object obj = py::reinterpret_borrow<py::object>(t);
if (checkMutableFunctionDefault(obj)) {
return true;
}
}
}
return false;
}
void checkMutableFunctionDefault(
const SourceRange& range,
const Argument& arg,
const py::object& def_arg) {
if (checkMutableFunctionDefault(def_arg) || arg.type()->cast<ClassType>()) {
throw ErrorReport(range)
<< "Mutable default parameters are not supported because Python binds them to the function"
<< " and they persist across function calls.\n As a workaround, make the default None and instantiate"
<< " the default parameter within the body of the function. Found "
<< def_arg.get_type() << " on parameter " << arg.name();
}
}
FunctionSchema getSchemaWithNameAndDefaults(
const SourceRange& range,
const FunctionSchema& schema,
const at::optional<std::string>& new_name,
const FunctionDefaults& default_args) {
std::vector<Argument> new_args;
for (auto& arg : schema.arguments()) {
auto it = default_args.find(arg.name());
if (it != default_args.end()) {
checkMutableFunctionDefault(range, arg, it->second);
c10::optional<IValue> value = tryCalculateDefaultParam(arg, it->second);
if (!value) {
ErrorReport error(range);
error << "Expected a default value of type " << arg.type()->repr_str()
<< " on parameter \"" << arg.name() << "\".";
if (arg.is_inferred_type()) {
error << "Because \"" << arg.name()
<< "\" was not annotated with an explicit type "
<< "it is assumed to be type 'Tensor'.";
}
throw error;
}
new_args.emplace_back(
arg.name(), arg.type(), arg.N(), *value, arg.kwarg_only());
} else {
new_args.push_back(arg);
}
}
return FunctionSchema(
new_name.value_or(schema.name()),
schema.overload_name(),
new_args,
schema.returns(),
schema.is_vararg(),
schema.is_varret());
}
static Decl mergeDefaultsAndExtraParametersToOverloadDecl(
const Decl& overload_decl,
const Decl& impl_decl,
const FunctionDefaults& defaults) {
std::vector<Param> adjusted_params;
const auto& overload_params = overload_decl.params();
const auto& impl_params = impl_decl.params();
// following PEP specification that the following should work:
// @overload
// def mouse_event(x1: int, y1: int) -> ClickEvent: ...
// ...
// def mouse_event(x1: int, y1: int, x2: Optional[int] = None, y2:
// Optional[int] = None)
TORCH_CHECK(
overload_params.size() <= impl_params.size(),
"Overload should not have more parameters than implementation function",
overload_decl.range(),
impl_decl.range());
for (const auto i : c10::irange(overload_params.size())) {
auto overload_name = overload_params[i].ident().name();
auto impl_name = impl_params[i].ident().name();
if (overload_name != impl_name) {
throw ErrorReport(overload_decl.range())
<< "Overload parameters must have the same names. "
<< "Found " << overload_name << " and " << impl_name
<< " on argument " << i;
}
adjusted_params.push_back(overload_params[i]);
}
for (size_t i = overload_params.size(); i < impl_params.size(); ++i) {
if (!defaults.count(impl_params[i].ident().name())) {
throw ErrorReport(impl_decl.range())
<< "Expected to find default parameter on argument"
<< impl_params[i].ident().name()
<< " because it is not defined on the overloaded declaration";
}
if (!impl_params[i].type().present()) {
throw ErrorReport(impl_decl.range())
<< "Parameters not specified on the overloaded declaration must have a type annotation in the implementation function."
<< " Did not find type for param " << impl_params[i].ident().name();
}
adjusted_params.push_back(impl_params[i]);
}
return Decl::create(
overload_decl.range(),
List<Param>::create(overload_decl.range(), adjusted_params),
overload_decl.return_type());
}
static StrongFunctionPtr script_compile_overloaded_function(
const c10::QualifiedName& name,
const Decl& overload_decl,
const Def& implementation_def,
const ResolutionCallback& rcb,
const FunctionDefaults& implementation_defaults,
const py::object& signature) {
if (signature.is(py::none())) {
throw ErrorReport(overload_decl.range())
<< "Must explicitly add type annotations to overloaded functions";
}
auto adjusted_decl = mergeDefaultsAndExtraParametersToOverloadDecl(
overload_decl, implementation_def.decl(), implementation_defaults);
auto new_def = implementation_def.withDecl(adjusted_decl);
auto cu = get_python_cu();
auto defined_functions = cu->define(
QualifiedName(name.prefix()),
/*properties=*/{},
/*propResolvers=*/{},
{new_def},
{pythonResolver(rcb)},
nullptr,
true);
TORCH_INTERNAL_ASSERT(defined_functions.size() == 1);
auto& defined = defined_functions[0];
FunctionDefaults updated_defaults = calcOverloadedFunctionDefaults(
defined->getSchema(), implementation_defaults);
defined->setSchema(getSchemaWithNameAndDefaults(
new_def.range(),
defined->getSchema(),
new_def.name().name(),
updated_defaults));
StrongFunctionPtr ret(std::move(cu), defined);
didFinishEmitFunction(ret);
return ret;
}
static StrongFunctionPtr script_compile_function(
const c10::QualifiedName& name,
const Def& def,
const FunctionDefaults& defaults,
const ResolutionCallback& rcb) {
auto cu = get_python_cu();
auto defined_functions = cu->define(
QualifiedName(name.prefix()),
/*properties=*/{},
/*propResolvers=*/{},
{def},
{pythonResolver(rcb)},
nullptr,
true);
TORCH_INTERNAL_ASSERT(defined_functions.size() == 1);
auto& defined = defined_functions[0];
defined->setSchema(getSchemaWithNameAndDefaults(
def.range(), defined->getSchema(), def.name().name(), defaults));
StrongFunctionPtr ret(std::move(cu), defined);
didFinishEmitFunction(ret);
return ret;
}
struct VISIBILITY_HIDDEN ModuleSelf : public Self {
ModuleSelf(std::shared_ptr<ConcreteModuleType> concreteType)
: Self(), concreteType_(std::move(concreteType)) {}
std::shared_ptr<SugaredValue> makeSugared(Value* v) const override {
v->setType(getClassType());
return std::make_shared<ModuleValue>(v, concreteType_);
}
ClassTypePtr getClassType() const override {
return concreteType_->getJitType()->expect<ClassType>();
}
private:
std::shared_ptr<ConcreteModuleType> concreteType_;
};
static TypePtr getTensorType(const at::Tensor& t, bool complete) {
auto r = TensorType::create(t);
if (!complete) {
r = r->dimensionedOnly();
}
return r;
}
static TypePtr inferShapeAndTypeForInput(
TypePtr input_type,
Stack::const_iterator& s_iter,
const Stack::const_iterator& s_iter_end,
bool complete);
static TupleTypePtr getTupleTensorType(
Stack::const_iterator& s_iter,
const Stack::const_iterator& s_iter_end,
const TypePtr& tupleType,
bool complete) {
TORCH_INTERNAL_ASSERT(tupleType->kind() == TupleType::Kind);
std::vector<TypePtr> types;
for (const auto& subType : tupleType->containedTypes()) {
TORCH_INTERNAL_ASSERT(s_iter != s_iter_end);
types.emplace_back(
inferShapeAndTypeForInput(subType, s_iter, s_iter_end, complete));
}
return TupleType::create(types);
}
static TypePtr inferShapeAndTypeForInput(
TypePtr input_type,
Stack::const_iterator& s_iter,
const Stack::const_iterator& s_iter_end,
bool complete) {
if (input_type->kind() == TupleType::Kind) {
return getTupleTensorType(s_iter, s_iter_end, input_type, complete);
} else if (input_type->kind() == TensorType::Kind) {
auto type = getTensorType(s_iter->toTensor(), complete);
s_iter++;
return type;
} else {
// Primitive type, keep as is.
s_iter++;
return input_type;
}
}
static void setInputTensorTypes(
Graph& g,
const Stack& stack,
bool complete,
const std::vector<int>& param_count_list = {}) {
at::ArrayRef<Value*> input_values = g.inputs();
auto s_iter = stack.begin();
size_t list_idx = 0;
if (!param_count_list.empty()) {
TORCH_INTERNAL_ASSERT(input_values.size() == param_count_list.size());
}
for (auto v : input_values) {
AT_ASSERT(s_iter != stack.end());
// Leave packed param types alone. This is needed for downstream passes
// (like alias analysis) to work properly. This will be unpacked later
// in unpackQuantizedWeights.
if (auto named_type = v->type()->cast<c10::NamedType>()) {
if (auto qualname = named_type->name()) {
if (getCustomClass(qualname->qualifiedName())) {
if (param_count_list.empty()) {
s_iter++;
} else {
s_iter += param_count_list[list_idx];
}
list_idx++;
continue;
}
}
}
v->setType(
inferShapeAndTypeForInput(v->type(), s_iter, stack.end(), complete));
list_idx++;
}
}
static std::shared_ptr<Graph> _propagate_shapes(
Graph& graph,
std::vector<at::Tensor> inputs,
bool with_grad = false) {
Stack stack(inputs.begin(), inputs.end());
auto retval = graph.copy();
setInputTensorTypes(*retval, stack, /*complete=*/false);
PropagateInputShapes(retval);
return retval;
}
static std::shared_ptr<Graph> _propagate_and_assign_input_shapes(
Graph& graph,
const std::vector<at::Tensor>& inputs,
const std::vector<int>& param_count_list,
bool with_grad = false,
bool propagate = true) {
auto retval = graph.copy();
setInputTensorTypes(
*retval, fmap<IValue>(inputs), /*complete=*/true, param_count_list);
if (propagate) {
PropagateInputShapes(retval);
}
return retval;
}
void addFunctionToModule(Module& module, const StrongFunctionPtr& func) {
// Make a graph with a fake self argument
auto graph = toGraphFunction(*func.function_).graph()->copy();
auto v = graph->insertInput(0, "self");
v->setType(module._ivalue()->type());
const auto name = QualifiedName(*module.type()->name(), "forward");
auto method =
module._ivalue()->compilation_unit()->create_function(name, graph);
module.type()->addMethod(method);
}
// this is used in our test suite to check that we correctly preserved type tags
bool ivalue_tags_match(const Module& lhs, const Module& rhs) {
struct Work {
IValue a;
IValue b;
};
std::unordered_set<const void*> visited;
std::vector<Work> work = {{lhs._ivalue(), rhs._ivalue()}};
while (!work.empty()) {
Work item = work.back();
work.pop_back();
if (item.a.isPtrType()) {
// uncomment to debug type matching errors
// std::cout << "MATCHING " << /*item.a <<*/ "(" << *item.a.type() << ") "
// << item.a.internalToPointer() << " " << /*item.b <<*/ " ("
// << *item.b.type() << ") " << item.b.internalToPointer() <<
// "\n";
if (visited.count(item.a.internalToPointer())) {
continue;
}
visited.emplace(item.a.internalToPointer());
}
if (!unshapedType(item.b.type())
->isSubtypeOf(unshapedType(item.b.type()))) {
// Since named types are saved and loaded in the test suite, we cannot
// expect them to be equal. We should still check their slots however.
if (!item.a.type()->cast<c10::NamedType>()) {
return false;
}
}
// check tags for objects that contain subobjects
if (item.a.isObject()) {
auto ao = item.a.toObject();
auto bo = item.b.toObject();
for (size_t i = 0; i < ao->slots().size(); ++i) {
work.emplace_back(Work{ao->slots().at(i), bo->slots().at(i)});
}
} else if (item.a.isTuple()) {
auto at = item.a.toTuple();
auto bt = item.b.toTuple();
for (size_t i = 0; i < at->elements().size(); ++i) {
work.emplace_back(Work{at->elements().at(i), bt->elements().at(i)});
}
} else if (item.a.isList()) {
auto al = item.a.toList();
auto bl = item.b.toList();
for (const auto i : c10::irange(al.size())) {
work.emplace_back(Work{al.get(i), bl.get(i)});
}
} else if (item.a.isGenericDict()) {
auto ad = item.a.toGenericDict();
auto bd = item.b.toGenericDict();
for (auto& item : ad) {
// Dictionaory keys cannot contain List/Dicts that require tags
// so we do not have to check them.
// Furthermore without ordered dicts it is expensive to find the
// equivalent key
work.emplace_back(Work{item.value(), bd.at(item.key())});
}
} else if (item.a.isFuture()) {
auto af = item.a.toFuture();
auto bf = item.b.toFuture();
af->wait();
bf->wait();
work.emplace_back(Work{af->value(), bf->value()});
}
}
return true;
}
// helper used to implement ._parameters, ._buffers, ._modules dicts
// inside of script nn.Module
template <typename Policy>
struct slot_dict_impl {
slot_dict_impl(ModulePtr module) : module_(std::move(module)) {}
bool contains(const std::string& name) const {
if (auto slot = module_->type()->findAttributeSlot(name)) {
if (Policy::valid(module_->type(), *slot, module_->getSlot(*slot))) {
return true;
}
}
return false;
}
std::vector<std::pair<std::string, py::object>> items() const {
std::vector<std::pair<std::string, py::object>> result;
for (size_t i = 0, N = module_->type()->numAttributes(); i < N; ++i) {
if (Policy::valid(module_->type(), i, module_->getSlot(i))) {
result.emplace_back(
module_->type()->getAttributeName(i),
toPyObject(module_->getSlot(i)));
}
}
return result;
}
void setattr(const std::string& name, py::object value) {
const TypePtr& type = module_->type()->getAttribute(name);
Module(module_).setattr(name, toIValue(std::move(value), type));
}
py::object getattr(const std::string& name) {
return toPyObject(Module(module_).attr(name));
}
static void bind(const py::module& m, const char* name) {
py::class_<slot_dict_impl<Policy>>(m, name)
.def(py::init(
[](Module& m) { return slot_dict_impl<Policy>(m._ivalue()); }))
.def("contains", &slot_dict_impl<Policy>::contains)
.def("items", &slot_dict_impl<Policy>::items)
.def("setattr", &slot_dict_impl<Policy>::setattr)
.def("getattr", &slot_dict_impl<Policy>::getattr);
}
private:
ModulePtr module_;
};
template <typename T>
py::list debugMakeList(const T& list) {
py::list result;
for (const auto& elem : list) {
result.append(py::cast(elem));
}
return result;
}
template <typename T>
py::list debugMakeNamedList(const T& list) {
py::list result;
for (auto elem : list) {
result.append(py::cast(std::make_pair(elem.name, elem.value)));
}
return result;
}
template <typename T>
py::set debugMakeSet(const T& list) {
py::set result;
for (const auto& elem : list) {
result.add(py::cast(elem));
}
return result;
}
static py::dict _jit_debug_module_iterators(Module& module) {
py::dict result;
result["children"] = debugMakeList(module.children());
result["named_children"] = debugMakeNamedList(module.named_children());
result["modules"] = debugMakeList(module.modules());
result["named_modules"] = debugMakeNamedList(module.named_modules());
result["parameters"] = debugMakeList(module.parameters(false));
result["named_parameters"] =
debugMakeNamedList(module.named_parameters(false));
result["parameters_r"] = debugMakeList(module.parameters(true));
result["named_parameters_r"] =
debugMakeNamedList(module.named_parameters(true));
result["buffers"] = debugMakeList(module.buffers(false));
result["named_buffers"] = debugMakeNamedList(module.named_buffers(false));
result["buffers_r"] = debugMakeList(module.buffers(true));
result["named_buffers_r"] = debugMakeNamedList(module.named_buffers(true));
result["named_attributes"] =
debugMakeNamedList(module.named_attributes(false));
result["named_attributes_r"] =
debugMakeNamedList(module.named_attributes(true));
return result;
}
static constexpr std::array<const char*, 47> magic_method_names = {
"__lt__", "__le__", "__eq__", "__ne__",
"__ge__", "__gt__", "__not__", "__abs__",
"__add__", "__and__", "__floordiv__", "__index__",
"__inv__", "__invert__", "__lshift__", "__mod__",
"__mul__", "__matmul__", "__neg__", "__or__",
"__pos__", "__pow__", "__rshift__", "__sub__",
"__truediv__", "__xor__", "__concat__", "__contains__",
"__delitem__", "__getitem__", "__setitem__", "__iadd__",
"__iand__", "__iconcat__", "__ifloordiv__", "__ilshift__",
"__imod__", "__imul__", "__imatmul__", "__ior__",
"__ipow__", "__irshift__", "__isub__", "__itruediv__",
"__ixor__", "__str__", "__len__",
};
struct DeepCopyMemoTable {
std::shared_ptr<IValue::HashAliasedIValueMap> map;
};
IValue pyIValueDeepcopy(const IValue& ivalue, const py::dict& memo) {
if (!memo.contains(py::str("__torch_script_memo_table"))) {
memo["__torch_script_memo_table"] =
DeepCopyMemoTable{std::make_shared<IValue::HashAliasedIValueMap>()};
}
auto& ivalue_memo =
*py::cast<DeepCopyMemoTable>(memo["__torch_script_memo_table"]).map;
return ivalue.deepcopy(ivalue_memo);
}
ExtraFilesMap extra_files_from_python(const py::dict& pydict) {
ExtraFilesMap r;
for (const auto& it : pydict) {
r[py::cast<std::string>(it.first)] = "";
}
return r;
}
void extra_files_to_python(const ExtraFilesMap& m, const py::dict& pydict) {
// py::dict is pointer-like type so it gets modified despite const&
for (const auto& it : m) {
pydict[py::str(it.first)] = py::bytes(it.second);
}
}
void pyCompilationUnitDefine(
CompilationUnit& cu,
const std::string& src,
const ResolutionCallback* rcb,
const uint32_t _frames_up) {
if (rcb && *rcb) {
cu.define(c10::nullopt, src, pythonResolver(*rcb), nullptr);
} else {
py::object py_default_rcb =
py::module::import("torch._jit_internal")
.attr("createResolutionCallbackFromFrame")(_frames_up);
auto default_rcb = py_default_rcb.cast<ResolutionCallback>();
cu.define(c10::nullopt, src, pythonResolver(default_rcb), nullptr);
}
}
void initJitScriptBindings(PyObject* module) {
auto m = py::handle(module).cast<py::module>();
// NOLINTNEXTLINE(bugprone-unused-raii)
py::class_<c10::Capsule>(m, "Capsule");
auto object_class =
py::class_<Object>(m, "ScriptObject")
.def("_type", [](Module& m) { return m.type(); })
.def(
"_get_method",
[](Object& self, const std::string& name) -> Method {
return self.get_method(name);
},
py::keep_alive<0, 1>())
.def(
"setattr",
[](Object& self, const std::string& name, py::object value) {
if (self.type()->hasConstant(name)) {
TORCH_CHECK(
false,
"Can't set constant '",
name,
"' which has value:",
self.type()->getConstant(name));
}
TypePtr type = self.type()->getAttribute(name);
try {
auto ivalue = toIValue(std::move(value), type);
self.setattr(name, ivalue);
} catch (std::exception& e) {
throw py::cast_error(c10::str(
"Could not cast attribute '",
name,
"' to type ",
type->repr_str(),
": ",
e.what()));
}
})
.def(
"getattr",
[](Object& self, const std::string& name) {
try {
return toPyObject(self.attr(name));
} catch (const ObjectAttributeError& err) {
throw AttributeError("%s", err.what());
}
})
.def(
"__getattr__",
[](Object& self, const std::string& name) -> py::object {
try {
if (name == "__qualname__") {
return py::cast(self.type()->name()->name());
}
if (auto method = self.find_method(name)) {
return py::cast(*method);
}
if (self.has_property(name)) {
auto prop = self.get_property(name);
// wrap the Method into callable PyObject
auto getter_func = py::cast(prop.getter_func);
return getter_func();
}
return toPyObject(self.attr(name));
} catch (const ObjectAttributeError& err) {
throw AttributeError("%s", err.what());
}
})
.def(
"__setattr__",
[](Object& self, const std::string& name, py::object value) {
try {
if (self.has_property(name)) {
auto prop = self.get_property(name);
if (!prop.setter_func.has_value()) {
TORCH_CHECK(false, "can't set attribute");
}
// wrap the Method into callable PyObject
auto setter_func = py::cast(prop.setter_func);
setter_func(value);
return;
}
if (self.type()->hasConstant(name)) {
TORCH_CHECK(
false,
"Can't set constant '",
name,
"' which has value:",
self.type()->getConstant(name));
}
TypePtr type = self.type()->getAttribute(name);
auto ivalue = toIValue(std::move(value), type);
self.setattr(name, ivalue);
} catch (const ObjectAttributeError& err) {
throw AttributeError("%s", err.what());
}
})
.def(
"hasattr",
[](Object& self, const std::string& name) {
return self.hasattr(name);
})
.def(
"_has_method",
[](Object& self, const std::string& name) {
return bool(self.find_method(name));
})
.def(
"_method_names",
[](Object& self) {
return fmap(self.get_methods(), [](const Method& method) {
return method.name();
});
})
.def(
"_properties", [](Object& self) { return self.get_properties(); })
.def("__copy__", &Object::copy)
.def(
"__hash__",
[](const Object& self) {
// Similar to Tensor's `__hash__`, which is `id()`.
return std::hash<c10::ivalue::Object*>{}(self._ivalue().get());
})
.def(py::pickle(
[](const Object& self)
-> std::tuple<py::object, std::string> { // __getstate__
if (auto getstate_method = self.find_method("__getstate__")) {
auto object_state = toPyObject((*getstate_method)(Stack{}));
TORCH_INTERNAL_ASSERT(self.type()->name());
return std::make_tuple(
object_state, self.type()->name()->qualifiedName());
}
std::stringstream err;
err << "Tried to serialize object ";
if (auto qualname = self.type()->name()) {
err << qualname->qualifiedName() << " ";
}
err << "which does not have a __getstate__ method defined!";
throw std::runtime_error(err.str());
},
[](const std::tuple<py::object, std::string>& state_tup)
-> Object {
py::object state;
std::string qualname;
std::tie(state, qualname) = state_tup;
auto class_type = getCustomClass(qualname);
TORCH_CHECK(
class_type,
"Tried to deserialize class ",
qualname,
" which is not known to the runtime. "
"If this is a custom C++ class, make "
"sure the appropriate code is linked.");
auto self = Object(c10::ivalue::Object::create(
c10::StrongTypePtr(
std::shared_ptr<torch::jit::CompilationUnit>(),
class_type),
1));
if (auto setstate_method = self.find_method("__setstate__")) {
auto setstate_schema =
setstate_method->function().getSchema();
TORCH_INTERNAL_ASSERT(
setstate_schema.arguments().size() == 2,
"__setstate__ method for class ",
class_type->repr_str(),
" must have exactly 2 arguments!");
auto state_type = setstate_schema.arguments().at(1).type();
(*setstate_method)(Stack{toIValue(state, state_type)});
return self;
}
std::stringstream err;
err << "Tried to deserialize object ";
if (auto qualname = class_type->name()) {
err << qualname->qualifiedName() << " ";
}
err << "which does not have a __setstate__ method defined!";
throw std::runtime_error(err.str());
}));
py::class_<Object::Property>(m, "ScriptObjectProperty")
.def_property_readonly(
"name", [](const Object::Property& self) { return self.name; })
.def_property_readonly(
"getter",
[](const Object::Property& self) { return self.getter_func; })
.def_property_readonly("setter", [](const Object::Property& self) {
return self.setter_func;
});
// Special case __str__ to make sure we can print Objects/Modules
// regardless of if the user defined a __str__
using MagicMethodImplType = std::function<py::object(
const Object& self, py::args args, py::kwargs kwargs)>;
std::unordered_map<std::string, MagicMethodImplType> special_magic_methods{
{"__str__",
[](const Object& self, py::args args, py::kwargs kwargs) -> py::object {
auto method = self.find_method("__str__");
if (!method) {
return py::str("ScriptObject");
}
return invokeScriptMethodFromPython(
*method,
// NOLINTNEXTLINE(performance-move-const-arg)
std::move(args),
// NOLINTNEXTLINE(performance-move-const-arg)
std::move(kwargs));
}}};
for (const char* mm_name : magic_method_names) {
if (special_magic_methods.count(mm_name)) {
object_class.def(mm_name, special_magic_methods[mm_name]);
} else {
object_class.def(
mm_name,
[mm_name](const Object& self, py::args args, py::kwargs kwargs) {
auto method = self.find_method(mm_name);
if (!method) {
throw NotImplementedError();
}
return invokeScriptMethodFromPython(
*method,
// NOLINTNEXTLINE(performance-move-const-arg)
std::move(args),
// NOLINTNEXTLINE(performance-move-const-arg)
std::move(kwargs));
});
}
}
// NOLINTNEXTLINE(bugprone-unused-raii)
py::class_<DeepCopyMemoTable>(m, "DeepCopyMemoTable");
py::class_<UpgraderEntry>(m, "_UpgraderEntry")
.def(py::init<int, std::string, std::string>())
.def_property_readonly(
"bumped_at_version",
[](const UpgraderEntry& self) { return self.bumped_at_version; })
.def_property_readonly(
"upgrader_name",
[](const UpgraderEntry& self) { return self.upgrader_name; })
.def_property_readonly("old_schema", [](const UpgraderEntry& self) {
return self.old_schema;
});
py::class_<UpgraderRange>(m, "_UpgraderRange")
.def(py::init<int, int>())
.def_property_readonly(
"min_version",
[](const UpgraderRange& self) { return self.min_version; })
.def_property_readonly("max_version", [](const UpgraderRange& self) {
return self.max_version;
});
object_class.def(
"__deepcopy__", [](const Object& self, const py::dict& memo) {
return Object(
pyIValueDeepcopy(IValue(self._ivalue()), memo).toObject());
});
// Used by torch.package to save ScriptModule objects in unified format.
py::class_<ScriptModuleSerializer>(m, "ScriptModuleSerializer")
.def(py::init<caffe2::serialize::PyTorchStreamWriter&>())
.def("serialize", &ScriptModuleSerializer::serialize_unified_format)
.def(
"write_files",
&ScriptModuleSerializer::writeFiles,
py::arg("code_dir") = ".data/ts_code/code/")
.def(
"storage_context",
&ScriptModuleSerializer::storage_context,
pybind11::return_value_policy::reference_internal);
// Used by torch.package to coordinate sharing of storages between eager
// and ScriptModules.
py::class_<
SerializationStorageContext,
std::shared_ptr<SerializationStorageContext>>(
m, "SerializationStorageContext")
.def("has_storage", &SerializationStorageContext::hasStorage)
.def("get_or_add_storage", &SerializationStorageContext::getOrAddStorage);
// torch.jit.ScriptModule is a subclass of this C++ object.
// Methods here are prefixed with _ since they should not be
// public.
py::class_<Module, Object>(m, "ScriptModule")
.def(py::init<std::string, std::shared_ptr<CompilationUnit>, bool>())
.def(
"save",
[](Module& m,
const std::string& filename,
const ExtraFilesMap& _extra_files = ExtraFilesMap()) {
m.save(filename, _extra_files);
},
py::arg("filename"),
py::arg("_extra_files") = ExtraFilesMap())
.def(
"save_to_buffer",
[](Module& m, const ExtraFilesMap& _extra_files = ExtraFilesMap()) {
std::ostringstream buf;
m.save(buf, _extra_files);
return py::bytes(buf.str());
},
py::arg("_extra_files") = ExtraFilesMap())
.def(
"_save_for_mobile",
[](Module& m,
const std::string& filename,
const ExtraFilesMap& _extra_files = ExtraFilesMap(),
bool _save_mobile_debug_info = false,
bool _use_flatbuffer = false) {
m._save_for_mobile(
filename,
_extra_files,
_save_mobile_debug_info,
_use_flatbuffer);
},
py::arg("filename"),
py::arg("_extra_files") = ExtraFilesMap(),
py::arg("_save_mobile_debug_info") = false,
py::arg("_use_flatbuffer") = false)
.def(
"_save_to_buffer_for_mobile",
[](Module& m,
const ExtraFilesMap& _extra_files = ExtraFilesMap(),
bool _save_mobile_debug_info = false,
bool _use_flatbuffer = false) {
std::ostringstream buf;
m._save_for_mobile(
buf, _extra_files, _save_mobile_debug_info, _use_flatbuffer);
return py::bytes(buf.str());
},
py::arg("_extra_files") = ExtraFilesMap(),
py::arg("_save_mobile_debug_info") = false,
py::arg("_use_flatbuffer") = false)
.def("_set_optimized", &Module::set_optimized)
.def(
"dump",
&Module::dump,
py::arg("code") = true,
py::arg("attrs") = true,
py::arg("params") = true)
.def(
"dump_to_str",
&Module::dump_to_str,
py::arg("code") = true,
py::arg("attrs") = true,
py::arg("params") = true)
.def(
"_replicate_for_data_parallel",
[](Module& module) {
const ModulePtr& obj = module._ivalue();
auto copy = c10::ivalue::Object::create(
c10::StrongTypePtr(obj->compilation_unit(), obj->type()),
obj->slots().size());
for (size_t i = 0; i < obj->slots().size(); ++i) {
copy->setSlot(i, obj->getSlot(i));
}
return Module(std::move(copy));
})
.def(
"get_debug_state",
[](Module& self) {
if (auto m = self.find_method("forward")) {
return m->get_executor().getDebugState();
}
throw std::runtime_error(
"Attempted to call get_debug_state on a Module without a compiled forward()");
})
.def(
"_define",
[](Module& m,
std::shared_ptr<ConcreteModuleType> concreteType,
const std::string& script,
const ResolutionCallback& rcb) {
const auto self = ModuleSelf(std::move(concreteType));
m._ivalue()->compilation_unit()->define(
*m.type()->name(), script, pythonResolver(rcb), &self);
didFinishEmitModule(m);
})
.def(
"_register_attribute",
[](Module& m,
const std::string& name,
const TypePtr& type,
py::handle value) {
m.register_attribute(name, type, toIValue(value, type));
})
.def(
"_create_method_from_trace",
[](Module& self,
const std::string& name,
const py::function& func,
const py::tuple& input_tuple,
const py::function& var_name_lookup_fn,
bool strict,
bool force_outplace,
const std::vector<std::string>& argument_names) {
// prereq: Module's buffers and parameters are unique
// this was ensured in python before calling this function
auto typed_inputs = toTraceableStack(input_tuple);
std::shared_ptr<Graph> graph =
std::get<0>(tracer::createGraphByTracing(
func,
typed_inputs,
var_name_lookup_fn,
strict,
force_outplace,
&self,
argument_names));
const auto method_name = QualifiedName(*self.type()->name(), name);
auto fn = self._ivalue()->compilation_unit()->create_function(
method_name, graph);
self.type()->addMethod(fn);
didFinishEmitModule(self);
},
py::arg("name"),
py::arg("func"),
py::arg("input_tuple"),
py::arg("var_name_lookup_fn"),
py::arg("strict"),
py::arg("force_outplace"),
py::arg("argument_names") = std::vector<std::string>())
.def(
"_get_forward_hooks",
[](const Module& m) {
std::vector<StrongFunctionPtr> funcs;
for (auto& hook : m.type()->getForwardHooks()) {
funcs.emplace_back(
StrongFunctionPtr(m.type()->compilation_unit(), hook));
}
return funcs;
})
.def(
"_get_forward_pre_hooks",
[](const Module& m) {
std::vector<StrongFunctionPtr> funcs;
for (auto& pre_hook : m.type()->getForwardPreHooks()) {
funcs.emplace_back(
StrongFunctionPtr(m.type()->compilation_unit(), pre_hook));
}
return funcs;
})
.def_property_readonly(
"code",
[](Module& self) {
std::vector<at::IValue> constants;
PrintDepsTable deps;
PythonPrint pp(constants, deps);
pp.printNamedType(self.type());
return pp.str();
})
.def_property_readonly(
"code_with_constants",
[](Module& self) {
std::vector<at::IValue> constants;
PrintDepsTable deps;
PythonPrint pp(constants, deps);
pp.printNamedType(self.type());
std::map<std::string, at::IValue> consts;
int i = 0;
for (auto const& constant : constants) {
consts["c" + std::to_string(i)] = constant;
i += 1;
}
return std::make_tuple(pp.str(), consts);
})
.def("apply", &Module::apply)
.def("__copy__", &Module::copy)
.def(
"__hash__",
[](const Module& self) {
// Similar to Tensor's `__hash__`, which is `id()`.
return std::hash<c10::ivalue::Object*>{}(self._ivalue().get());
})
.def(
"__eq__",
[](const Module& self, const py::object& other) {
// TODO: call UDF if it exists
if (!py::isinstance<Module>(other)) {
return false;
}
return self._ivalue().get() ==
py::cast<Module>(other)._ivalue().get();
})
.def(
"__deepcopy__",
[](const Module& self, const py::dict& memo) {
return Module(
pyIValueDeepcopy(IValue(self._ivalue()), memo).toObject());
})
.def("children", &Module::children)
.def_property_readonly("qualified_name", [](const Module& self) {
return self.type()->name()->qualifiedName();
});
py::class_<mobile::Module>(m, "LiteScriptModule")
.def(py::init<
c10::intrusive_ptr<c10::ivalue::Object>,
std::shared_ptr<mobile::CompilationUnit>>())
.def(
"find_method",
[](mobile::Module& m, const std::string& method_name) {
auto method = m.find_method(method_name);
return method != c10::nullopt;
},
py::arg("method_name"))
.def(
"run_method",
[](mobile::Module& m,
const std::string& method_name,
const py::tuple& input_tuple) {
Stack stack;
for (auto& input : input_tuple) {
stack.push_back(toTypeInferredIValue(input));
}
return m.get_method(method_name)(stack);
},
py::arg("method_name"),
py::arg("input_tuple"))
.def(
"forward",
[](mobile::Module& m, const py::tuple& input_tuple) {
Stack stack;
for (auto& input : input_tuple) {
stack.push_back(toTypeInferredIValue(input));
}
return m.get_method("forward")(stack);
},
py::arg("input_tuple"));
slot_dict_impl<detail::ParameterPolicy>::bind(m, "ParameterDict");
slot_dict_impl<detail::BufferPolicy>::bind(m, "BufferDict");
slot_dict_impl<detail::ModulePolicy>::bind(m, "ModuleDict");
py::class_<ErrorReport, std::shared_ptr<ErrorReport>>(m, "ErrorReport")
.def(py::init<SourceRange>())
.def("what", &ErrorReport::what)
.def_static("call_stack", ErrorReport::current_call_stack);
py::class_<CompilationUnit, std::shared_ptr<CompilationUnit>>(
m, "CompilationUnit")
.def(
py::init([](const std::string& lang, const uint32_t _frames_up) {
auto cu = std::make_shared<CompilationUnit>();
if (lang.size() > 0) {
pyCompilationUnitDefine(*cu, lang, nullptr, _frames_up);
}
return cu;
}),
py::arg("lang") = "",
py::arg("_frames_up") = 0)
.def(
"find_function",
[](std::shared_ptr<CompilationUnit> self, const std::string& name) {
auto fn = self->find_function(QualifiedName(name));
if (fn) {
return c10::optional<StrongFunctionPtr>(
StrongFunctionPtr(std::move(self), fn));
} else {
return c10::optional<StrongFunctionPtr>(c10::nullopt);
}
})
.def(
"__getattr__",
[](std::shared_ptr<CompilationUnit> self, const std::string& name) {
auto fn = self->find_function(QualifiedName(name));
if (fn) {
return StrongFunctionPtr(std::move(self), fn);
} else {
throw AttributeError(
"'CompilationUnit' has no attribute '%s'", name.c_str());
}
})
.def(
"get_functions",
[](const std::shared_ptr<CompilationUnit>& self) {
auto raw_functions = self->get_functions();
std::vector<StrongFunctionPtr> functions;
functions.reserve(raw_functions.size());
for (auto fn : raw_functions) {
if (fn) {
functions.emplace_back(self, fn);
}
}
return functions;
})
.def("set_optimized", &CompilationUnit::set_optimized)
.def(
"define",
pyCompilationUnitDefine,
py::arg("src"),
py::arg("rcb") = nullptr,
py::arg("_frames_up") = 0)
.def(
"create_function",
[](std::shared_ptr<CompilationUnit>& self,
const std::string& qualified_name,
std::shared_ptr<Graph> graph,
bool should_mangle) {
Function* fn = self->create_function(
qualified_name, std::move(graph), should_mangle);
return StrongFunctionPtr(std::move(self), fn);
},
py::arg("qualified_name"),
py::arg("graph"),
py::arg("should_mangle") = false)
.def(
"get_interface",
[](const std::shared_ptr<CompilationUnit>& self,
const std::string& name) { return self->get_interface(name); })
.def(
"get_class",
[](const std::shared_ptr<CompilationUnit>& self,
const std::string& name) { return self->get_class(name); })
.def(
"drop_all_functions",
[](const std::shared_ptr<CompilationUnit>& self) {
self->drop_all_functions();
});
py::class_<StrongFunctionPtr>(m, "ScriptFunction", py::dynamic_attr())
.def(
"__call__",
[](py::args args, py::kwargs kwargs) {
HANDLE_TH_ERRORS
// see: [pybind11 varargs]
auto strongPtr = py::cast<StrongFunctionPtr>(args[0]);
Function& callee = *strongPtr.function_;
py::object result = invokeScriptFunctionFromPython(
callee,
// NOLINTNEXTLINE(performance-move-const-arg)
tuple_slice(std::move(args), 1),
// NOLINTNEXTLINE(performance-move-const-arg)
std::move(kwargs));
return result;
END_HANDLE_TH_ERRORS_PYBIND
})
.def(
"save",
[](const StrongFunctionPtr& self,
const std::string& filename,
const ExtraFilesMap& _extra_files = ExtraFilesMap()) {
Module module("__torch__.PlaceholderModule");
// [issue 27343]
// Modules have 'training' attributes by default, but due to
// https://github.com/pytorch/pytorch/issues/27343, functions end
// up having a training attribute when they are loaded. This adds
// a fake 'training' attribute that shouldn't be used, but prevents
// jitter on saving and loading. Once that issue is fixed this can
// be deleted.
module.register_attribute("training", BoolType::get(), true);
addFunctionToModule(module, self);
module.save(filename, _extra_files);
},
py::arg("filename"),
py::arg("_extra_files") = ExtraFilesMap())
.def(
"save_to_buffer",
[](const StrongFunctionPtr& self,
const ExtraFilesMap& _extra_files = ExtraFilesMap()) {
std::ostringstream buf;
Module module("__torch__.PlaceholderModule");
// see [issue 27343]
module.register_attribute("training", BoolType::get(), true);
addFunctionToModule(module, self);
module.save(buf, _extra_files);
return py::bytes(buf.str());
},
py::arg("_extra_files") = ExtraFilesMap())
.def_property_readonly(
"graph",
[](const StrongFunctionPtr& self) {
return toGraphFunction(*self.function_).graph();
})
.def_property_readonly(
"inlined_graph",
[](const StrongFunctionPtr& self) {
auto g = toGraphFunction(*self.function_).graph()->copy();
Inline(*g);
return g;
})
.def_property_readonly(
"schema",
[](const StrongFunctionPtr& self) {
return self.function_->getSchema();
})
.def_property_readonly(
"code",
[](const StrongFunctionPtr& self) {
std::vector<at::IValue> constants;
PrintDepsTable deps;
PythonPrint pp(constants, deps);
pp.printFunction(*self.function_);
return pp.str();
})
.def(
"get_debug_state",
[](const StrongFunctionPtr& self) {
return toGraphFunction(*self.function_)
.get_executor()
.getDebugState();
})
.def(
"_debug_flush_compilation_cache",
[](const StrongFunctionPtr& self) {
toGraphFunction(*self.function_)
.get_executor()
.debugFlushCompilationCache();
})
.def_property_readonly(
"name",
[](const StrongFunctionPtr& self) { return self.function_->name(); })
.def_property_readonly(
"qualified_name",
[](const StrongFunctionPtr& self) {
return self.function_->qualname().qualifiedName();
})
.def_property_readonly("__doc__", [](const StrongFunctionPtr& self) {
return self.function_->doc_string();
});
py::class_<Method>(m, "ScriptMethod", py::dynamic_attr())
.def(
"__call__",
[](py::args args, py::kwargs kwargs) {
// see: [pybind11 varargs]
HANDLE_TH_ERRORS
Method& method = py::cast<Method&>(args[0]);
return invokeScriptMethodFromPython(
method,
// NOLINTNEXTLINE(performance-move-const-arg)
tuple_slice(std::move(args), 1),
// NOLINTNEXTLINE(performance-move-const-arg)
std::move(kwargs));
END_HANDLE_TH_ERRORS_PYBIND
})
.def_property_readonly("graph", &Method::graph)
.def_property_readonly(
"inlined_graph",
[](const Method& self) {
auto g = toGraphFunction(self.function()).graph()->copy();
Inline(*g);
return g;
})
.def_property_readonly(
"schema", [](Method& m) { return m.function().getSchema(); })
.def_property_readonly("name", &Method::name)
.def_property_readonly(
"code",
[](Method& self) {
std::vector<at::IValue> constants;
PrintDepsTable deps;
PythonPrint pp(constants, deps);
pp.printMethod(self.function());
return pp.str();
})
.def(
"_debug_flush_compilation_cache",
[](Method& self) {
return self.get_executor().debugFlushCompilationCache();
})
.def_property_readonly(
"code_with_constants",
[](Method& self) {
std::vector<at::IValue> constants;
PrintDepsTable deps;
PythonPrint pp(constants, deps);
pp.printMethod(self.function());
std::map<std::string, at::IValue> consts;
int i = 0;
for (auto const& constant : constants) {
consts["c" + std::to_string(i)] = constant;
i += 1;
}
return std::make_tuple(pp.str(), consts);
})
.def_property_readonly("owner", &Method::owner);
m.def("_generate_upgraders_graph", &generate_upgraders_graph);
m.def(
"_compile_graph_to_code_table",
[](const std::string& name, const std::shared_ptr<Graph>& graph) {
CompilationOptions options;
GraphFunction jitFunc(name, graph, nullptr);
auto mobileFunc = convertJitFunctionToMobileFunction(jitFunc, options);
return convertMobileFunctionToCodeTable(*mobileFunc, options);
});
m.def(
"_jit_script_compile",
[](const std::string& qualname,
const Def& def,
const ResolutionCallback& rcb,
const FunctionDefaults& defaults) {
C10_LOG_API_USAGE_ONCE("torch.script.compile");
const auto name = c10::QualifiedName(qualname);
TORCH_INTERNAL_ASSERT(name.name() == def.name().name());
return script_compile_function(name, def, defaults, rcb);
});
m.def(
"_jit_script_compile_overload",
[](const std::string& qualname,
const Decl& overload_decl,
const Def& implementation_def,
const ResolutionCallback& rcb,
const FunctionDefaults& implementation_defaults,
const py::object& signature) {
const auto name = c10::QualifiedName(qualname);
return script_compile_overloaded_function(
name,
overload_decl,
implementation_def,
rcb,
implementation_defaults,
signature);
});
m.def(
"_replace_overloaded_method_decl",
[](const Decl& overload_decl,
const Def& implementation_def,
const std::string& new_name) {
checkOverloadDecl(overload_decl, implementation_def.decl());
return implementation_def.withDecl(overload_decl).withName(new_name);
});
m.def(
"_create_function_from_trace",
[](const std::string& qualname,
const py::function& func,
const py::tuple& input_tuple,
const py::function& var_name_lookup_fn,
bool strict,
bool force_outplace,
const std::vector<std::string>& argument_names) {
auto typed_inputs = toTraceableStack(input_tuple);
std::shared_ptr<Graph> graph = std::get<0>(tracer::createGraphByTracing(
func,
typed_inputs,
var_name_lookup_fn,
strict,
force_outplace,
/*self=*/nullptr,
argument_names));
auto cu = get_python_cu();
auto name = c10::QualifiedName(qualname);
auto result = cu->create_function(
std::move(name), std::move(graph), /*shouldMangle=*/true);
StrongFunctionPtr ret(std::move(cu), result);
didFinishEmitFunction(ret);
return ret;
},
py::arg("name"),
py::arg("func"),
py::arg("input_tuple"),
py::arg("var_name_lookup_fn"),
py::arg("strict"),
py::arg("force_outplace"),
py::arg("argument_names") = std::vector<std::string>());
m.def(
"_jit_script_class_compile",
[](const std::string& qualifiedName,
const ClassDef& classDef,
const ClassMethodDefaults& defaults,
const ResolutionCallback& rcb) {
C10_LOG_API_USAGE_ONCE("torch.script.class");
if (classDef.superclass().present()) {
throw ErrorReport(classDef.range())
<< "Torchscript does not support class inheritance.";
}
auto cu = get_python_cu();
auto classname = c10::QualifiedName(qualifiedName);
if (cu->get_type(classname) != nullptr) {
classname = cu->mangle(classname);
}
auto classType = ClassType::create(
classname,
cu,
/* is_module = */ false,
/* doc_string = */ "",
getUnresolvedClassAttributes(classDef));
cu->register_type(classType);
std::vector<ResolverPtr> methodRcbs, propRcbs;
std::vector<Def> methodDefs;
std::vector<Property> props;
for (const auto& def : classDef.body()) {
if (def.kind() != TK_DEF) {
throw ErrorReport(def.range())
<< "Currently class bodies can only contain method "
"definitions. File an issue on Github if you want "
"something else!";
}
methodDefs.emplace_back(Def(def));
methodRcbs.push_back(
pythonResolver(rcb, classDef.name().name(), classType));
}
// Gather definitions for property getters and setters as well as
// corresponding resolution callbacks.
if (classDef.properties().present()) {
for (const auto& prop : classDef.properties().get()) {
props.emplace_back(prop);
propRcbs.push_back(
pythonResolver(rcb, classDef.name().name(), classType));
}
}
const auto self = SimpleSelf(classType);
cu->define(classname, props, propRcbs, methodDefs, methodRcbs, &self);
// Stitch in default arguments for methods. Properties don't need to be
// considered since there is no way to invoke setters without passing in
// a value.
auto defs_it = methodDefs.begin();
while (defs_it != methodDefs.end()) {
auto def_name = (*defs_it).name().name();
// If the method is not in the defaults map, assume there are
// no default arguments for it.
auto default_it = defaults.find(def_name);
if (default_it == defaults.end()) {
continue;
}
const auto method_name =
QualifiedName(classname, (*defs_it).name().name());
auto& method = cu->get_function(method_name);
method.setSchema(getSchemaWithNameAndDefaults(
defs_it->range(),
method.getSchema(),
at::nullopt,
default_it->second));
++defs_it;
}
return classType;
});
m.def(
"_jit_script_interface_compile",
[](const std::string& qualifiedName,
const ClassDef& classDef,
const ResolutionCallback& rcb,
bool is_module) {
auto cu = get_python_cu();
auto className = c10::QualifiedName(qualifiedName);
if (cu->get_type(className) != nullptr) {
className = cu->mangle(className);
}
get_python_cu()->define_interface(
className, classDef, pythonResolver(rcb), is_module);
return className.qualifiedName();
});
py::class_<torch::jit::ErrorReport::CallStack>(
m, "CallStack", py::dynamic_attr())
.def(py::init<const std::string&, const SourceRange&>());
m.def("_parse_source_def", [](const std::string& src) {
Parser p(std::make_shared<Source>(src));
return Def(p.parseFunction(/*is_method=*/true));
});
m.def("parse_type_comment", [](const std::string& comment) {
Parser p(std::make_shared<Source>(comment));
return Decl(p.parseTypeComment());
});
m.def("_is_upgraders_enabled", &is_upgraders_enabled);
m.def("_get_upgraders_map_size", &get_upgraders_map_size);
m.def("_dump_upgraders_map", &dump_upgraders_map);
m.def("_test_only_populate_upgraders", &test_only_populate_upgraders);
m.def("_test_only_remove_upgraders", &test_only_remove_upgraders);
m.def("merge_type_from_type_comment", &mergeTypesFromTypeComment);
m.def("_get_max_operator_version", &getMaxOperatorVersion);
m.def("_get_operator_version_map", &get_operator_version_map);
m.def("_get_upgrader_ranges", &getUpgradersRangeForOp);
m.def("_test_only_add_entry_to_op_version_map", &test_only_add_entry);
m.def("_test_only_remove_entry_to_op_version_map", &test_only_remove_entry);
m.def(
"import_ir_module",
[](std::shared_ptr<CompilationUnit> cu,
const std::string& filename,
py::object map_location,
const py::dict& extra_files) {
c10::optional<at::Device> optional_device;
if (!map_location.is(py::none())) {
AT_ASSERT(THPDevice_Check(map_location.ptr()));
optional_device =
reinterpret_cast<THPDevice*>(map_location.ptr())->device;
}
ExtraFilesMap extra_files_map = extra_files_from_python(extra_files);
auto ret = import_ir_module(
std::move(cu), filename, optional_device, extra_files_map);
extra_files_to_python(extra_files_map, extra_files);
return ret;
});
m.def(
"_import_ir_module_from_package",
[](std::shared_ptr<CompilationUnit> cu,
std::shared_ptr<caffe2::serialize::PyTorchStreamReader> reader,
std::shared_ptr<torch::jit::DeserializationStorageContext>
storage_context,
py::object map_location,
std::string ts_id) {
c10::optional<at::Device> optional_device;
if (!map_location.is(py::none())) {
AT_ASSERT(THPDevice_Check(map_location.ptr()));
optional_device =
reinterpret_cast<THPDevice*>(map_location.ptr())->device;
}
return import_ir_module(
std::move(cu),
std::move(reader),
std::move(storage_context),
optional_device,
std::move(ts_id));
});
m.def(
"import_ir_module_from_buffer",
[](std::shared_ptr<CompilationUnit> cu,
const std::string& buffer,
py::object map_location,
const py::dict& extra_files) {
std::istringstream in(buffer);
c10::optional<at::Device> optional_device;
if (!map_location.is(py::none())) {
AT_ASSERT(THPDevice_Check(map_location.ptr()));
optional_device =
reinterpret_cast<THPDevice*>(map_location.ptr())->device;
}
ExtraFilesMap extra_files_map = extra_files_from_python(extra_files);
auto ret = import_ir_module(
std::move(cu), in, optional_device, extra_files_map);
extra_files_to_python(extra_files_map, extra_files);
return ret;
});
m.def(
"_load_for_lite_interpreter",
[](const std::string& filename, py::object map_location) {
c10::optional<at::Device> optional_device;
if (!map_location.is(py::none())) {
AT_ASSERT(THPDevice_Check(map_location.ptr()));
optional_device =
reinterpret_cast<THPDevice*>(map_location.ptr())->device;
}
return _load_for_mobile(filename, optional_device);
});
m.def(
"_load_for_lite_interpreter_from_buffer",
[](const std::string& buffer, py::object map_location) {
std::istringstream in(buffer);
c10::optional<at::Device> optional_device;
if (!map_location.is(py::none())) {
AT_ASSERT(THPDevice_Check(map_location.ptr()));
optional_device =
reinterpret_cast<THPDevice*>(map_location.ptr())->device;
}
return _load_for_mobile(in, optional_device);
});
m.def(
"_backport_for_mobile",
[](const std::string& filename_input,
const std::string& filename_output,
const int64_t version) {
return _backport_for_mobile(filename_input, filename_output, version);
});
m.def(
"_backport_for_mobile_from_buffer",
[](const std::string& buffer_input,
const std::string& filename_output,
const int64_t version) {
std::istringstream in(buffer_input);
return _backport_for_mobile(in, filename_output, version);
});
m.def(
"_backport_for_mobile_to_buffer",
[](const std::string& filename_input, const int64_t version) {
std::ostringstream buffer_output;
bool success =
_backport_for_mobile(filename_input, buffer_output, version);
return success ? py::bytes(buffer_output.str()) : py::bytes("");
});
m.def(
"_backport_for_mobile_from_buffer_to_buffer",
[](const std::string& buffer_input, const int64_t version) {
std::istringstream in(buffer_input);
std::ostringstream buffer_output;
bool success = _backport_for_mobile(in, buffer_output, version);
return success ? py::bytes(buffer_output.str()) : py::bytes("");
});
m.def("_get_model_bytecode_version", [](const std::string& filename) {
return _get_model_bytecode_version(filename);
});
m.def(
"_get_model_bytecode_version_from_buffer", [](const std::string& buffer) {
std::istringstream in(buffer);
return _get_model_bytecode_version(in);
});
m.def("_get_mobile_model_contained_types", [](const std::string& filename) {
return _get_mobile_model_contained_types(filename);
});
m.def(
"_get_mobile_model_contained_types_from_buffer",
[](const std::string& buffer) {
std::istringstream in(buffer);
return _get_mobile_model_contained_types(in);
});
m.def("_nn_module_to_mobile", [](const Module& module) {
CompilationOptions options;
return jitModuleToMobile(module, options);
});
py::class_<OperatorInfo>(m, "OperatorInfo")
.def_readonly("num_schema_args", &OperatorInfo::num_schema_args);
m.def("_get_model_ops_and_info", [](const std::string& filename) {
return _get_model_ops_and_info(filename);
});
m.def("_get_model_ops_and_info_from_buffer", [](const std::string& buffer) {
std::istringstream in(buffer);
return _get_model_ops_and_info(in);
});
m.def("_export_operator_list", [](torch::jit::mobile::Module& sm) {
return debugMakeSet(torch::jit::mobile::_export_operator_list(sm));
});
m.def("_jit_set_emit_hooks", setEmitHooks);
m.def("_jit_get_emit_hooks", getEmitHooks);
m.def("_jit_clear_class_registry", []() {
get_python_cu()->_clear_python_cu();
});
m.def(
"_debug_set_autodiff_subgraph_inlining",
debugSetAutodiffSubgraphInlining);
m.def("_debug_set_fusion_group_inlining", debugSetFusionGroupInlining);
m.def("_debug_get_fusion_group_inlining", getFusionGroupInlining);
m.def("_propagate_shapes", _propagate_shapes);
m.def(
"_propagate_and_assign_input_shapes", _propagate_and_assign_input_shapes);
m.def(
"_last_executed_optimized_graph",
[]() { return lastExecutedOptimizedGraph(); },
"Retrieve the optimized graph that was run the last time the graph executor ran on this thread");
m.def(
"_create_function_from_graph",
[](const std::string& qualname, std::shared_ptr<Graph> graph) {
// TODO this should go in the global Python CU
auto cu = std::make_shared<CompilationUnit>();
c10::QualifiedName name(qualname);
auto fn = cu->create_function(std::move(name), std::move(graph));
return StrongFunctionPtr(std::move(cu), fn);
});
m.def("_ivalue_tags_match", ivalue_tags_match);
m.def("_ivalue_debug_python_object", [](py::object py_obj) {
// convert to IValue first, IValue will incref via py::object
IValue pyobj_ivalue = toIValue(std::move(py_obj), PyObjectType::get());
// convert back to PyObject by borrowing the reference, which also
// incref, after the return of this function, IValue is out of scope
// which decref, so the return value is original refcount + 1
py::object ret = toPyObject(pyobj_ivalue);
return ret;
});
m.def("_jit_debug_module_iterators", _jit_debug_module_iterators);
py::class_<testing::FileCheck>(m, "FileCheck")
.def(py::init<>())
.def("check", &testing::FileCheck::check)
.def("check_not", &testing::FileCheck::check_not)
.def("check_same", &testing::FileCheck::check_same)
.def("check_next", &testing::FileCheck::check_next)
.def("check_count", &testing::FileCheck::check_count)
.def("check_dag", &testing::FileCheck::check_dag)
.def(
"check_source_highlighted",
&testing::FileCheck::check_source_highlighted)
.def(
"check_count",
[](testing::FileCheck& f,
const std::string& str,
size_t count,
bool exactly) { return f.check_count(str, count, exactly); },
"Check Count",
py::arg("str"),
py::arg("count"),
py::arg("exactly") = false)
.def(
"run",
[](testing::FileCheck& f, const std::string& str) {
return f.run(str);
})
.def(
"run", [](testing::FileCheck& f, const Graph& g) { return f.run(g); })
.def(
"run",
[](testing::FileCheck& f,
const std::string& input,
const std::string& output) { return f.run(input, output); },
"Run",
py::arg("checks_file"),
py::arg("test_file"))
.def(
"run",
[](testing::FileCheck& f, const std::string& input, const Graph& g) {
return f.run(input, g);
},
"Run",
py::arg("checks_file"),
py::arg("graph"));
m.def(
"_logging_set_logger",
[](logging::LoggerBase* logger) { return logging::setLogger(logger); },
py::return_value_policy::reference);
m.def("_set_graph_executor_optimize", [](bool optimize) {
setGraphExecutorOptimize(optimize);
});
m.def("_get_graph_executor_optimize", &torch::jit::getGraphExecutorOptimize);
m.def(
"_enable_mobile_interface_call_export",
&torch::jit::enableMobileInterfaceCallExport);
m.def("_create_module_with_type", [](const ClassTypePtr& type) {
return Module(get_python_cu(), type);
}).def("_create_object_with_type", [](const ClassTypePtr& type) {
return Object(get_python_cu(), type);
});
m.def("_export_opnames", [](Module& sm) {
return debugMakeList(torch::jit::export_opnames(sm));
});
py::class_<
ConcreteModuleTypeBuilder,
std::shared_ptr<ConcreteModuleTypeBuilder>>(
m, "ConcreteModuleTypeBuilder")
.def(py::init<py::object>())
.def(
"add_constant",
[](ConcreteModuleTypeBuilder& self,
std::string name,
py::object value) {
self.addConstant(std::move(name), std::move(value));
})
.def("add_attribute", &ConcreteModuleTypeBuilder::addAttribute)
.def(
"add_function_attribute",
&ConcreteModuleTypeBuilder::addFunctionAttribute)
.def(
"add_builtin_function",
&ConcreteModuleTypeBuilder::addBuiltinFunction)
.def("add_forward_hook", &ConcreteModuleTypeBuilder::addForwardHook)
.def(
"add_forward_pre_hook", &ConcreteModuleTypeBuilder::addForwardPreHook)
.def("add_module", &ConcreteModuleTypeBuilder::addModule)
.def("add_overload", &ConcreteModuleTypeBuilder::addOverload)
.def("set_poisoned", &ConcreteModuleTypeBuilder::setPoisoned)
.def(
"add_failed_attribute",
&ConcreteModuleTypeBuilder::addFailedAttribute)
.def(
"add_ignored_attribute",
&ConcreteModuleTypeBuilder::addIgnoredAttribute)
.def(
"add_ignored_attributes",
[](ConcreteModuleTypeBuilder& self,
const std::vector<std::string>& names) {
for (auto& name : names) {
self.addIgnoredAttribute(name);
}
})
.def(
"set_module_dict",
[](ConcreteModuleTypeBuilder& self) {
self.setIterableModuleKind(IterableModuleKind::DICT);
})
.def("build", &ConcreteModuleTypeBuilder::build)
.def(
"equals",
[](const ConcreteModuleTypeBuilder& self,
const ConcreteModuleTypeBuilder& other) {
return self.equals(other);
})
.def("set_module_list", [](ConcreteModuleTypeBuilder& self) {
self.setIterableModuleKind(IterableModuleKind::LIST);
});
py::class_<ConcreteModuleType, std::shared_ptr<ConcreteModuleType>>(
m, "ConcreteModuleType")
.def_property_readonly("py_class", &ConcreteModuleType::getPyClass)
.def_property_readonly("jit_type", &ConcreteModuleType::getJitType)
.def_static("from_jit_type", &ConcreteModuleType::fromJitType)
.def("get_constants", &ConcreteModuleType::getConstantsPy)
.def("get_attributes", &ConcreteModuleType::getAttributesPy)
.def("get_modules", &ConcreteModuleType::getModulesPy)
.def("dump", &ConcreteModuleType::dump)
.def("is_ignored_attribute", &ConcreteModuleType::isIgnoredAttribute)
.def(
"equals",
[](const ConcreteModuleType& self, const ConcreteModuleType& other) {
return self.equals(other);
})
.def(
"equals",
[](const ConcreteModuleType& self,
const ConcreteModuleTypeBuilder& other) {
return self.equals(other);
})
.def(
"_create_methods_and_properties",
[](std::shared_ptr<ConcreteModuleType> concreteType,
const std::vector<Property>& properties,
const std::vector<ResolutionCallback>& propertyRcbs,
const std::vector<Def>& methodDefs,
const std::vector<ResolutionCallback>& methodRcbs,
const std::vector<FunctionDefaults>& defaults) {
TORCH_INTERNAL_ASSERT(methodDefs.size() == methodRcbs.size());
TORCH_INTERNAL_ASSERT(properties.size() == propertyRcbs.size());
std::vector<ResolverPtr> methodResolvers, propertyResolvers;
methodResolvers.reserve(methodRcbs.size());
for (auto& callback : methodRcbs) {
methodResolvers.push_back(pythonResolver(callback));
}
propertyResolvers.reserve(propertyRcbs.size());
for (auto& callback : propertyRcbs) {
propertyResolvers.push_back(pythonResolver(callback));
}
const auto& selfType =
concreteType->getJitType()->expect<ClassType>();
const auto& prefix = selfType->name().value();
const auto self = ModuleSelf(std::move(concreteType));
auto cu = selfType->compilation_unit();
cu->define(
prefix,
properties,
propertyResolvers,
methodDefs,
methodResolvers,
&self);
// Stitch in default arguments for each Def if provided
auto defaults_it = defaults.begin();
auto defs_it = methodDefs.begin();
while (defs_it != methodDefs.end()) {
const auto method_name =
QualifiedName(prefix, (*defs_it).name().name());
auto& method = cu->get_function(method_name);
method.setSchema(getSchemaWithNameAndDefaults(
defs_it->range(),
method.getSchema(),
at::nullopt,
*defaults_it));
++defs_it;
++defaults_it;
}
})
.def(
"_create_hooks",
[](std::shared_ptr<ConcreteModuleType> concreteType,
const std::vector<Def>& hookDefs,
const std::vector<ResolutionCallback>& hookRcbs,
const std::vector<Def>& preHookDefs,
const std::vector<ResolutionCallback>& preHookRcbs) {
TORCH_INTERNAL_ASSERT(hookDefs.size() == hookRcbs.size());
TORCH_INTERNAL_ASSERT(preHookDefs.size() == preHookRcbs.size());
std::vector<ResolverPtr> hookResolvers, preHookResolvers;
hookResolvers.reserve(hookRcbs.size());
for (auto& callback : hookRcbs) {
hookResolvers.push_back(pythonResolver(callback));
}
preHookResolvers.reserve(preHookRcbs.size());
for (auto& callback : preHookRcbs) {
preHookResolvers.push_back(pythonResolver(callback));
}
const auto& selfType =
concreteType->getJitType()->expect<ClassType>();
const auto& prefix = selfType->name().value();
const auto self = ModuleSelf(std::move(concreteType));
auto cu = selfType->compilation_unit();
cu->define_hooks(
prefix,
hookDefs,
hookResolvers,
preHookDefs,
preHookResolvers,
&self);
});
m.def(
"_resolve_type",
[](const std::string& name,
const SourceRange& range,
const ResolutionCallback& rcb) {
return pythonResolver(rcb)->resolveType(name, range);
});
m.def(
"_resolve_type_from_object",
[](const py::object& obj,
const SourceRange& range,
const ResolutionCallback& rcb) {
return pythonResolver(rcb)->resolveTypeFromObject(obj, range);
});
m.def(
"_run_emit_module_hook", [](const Module& m) { didFinishEmitModule(m); });
// NOLINTNEXTLINE(bugprone-unused-raii)
py::class_<logging::LoggerBase, std::shared_ptr<logging::LoggerBase>>(
m, "LoggerBase");
py::enum_<logging::LockingLogger::AggregationType>(m, "AggregationType")
.value("SUM", logging::LockingLogger::AggregationType::SUM)
.value("AVG", logging::LockingLogger::AggregationType::AVG)
.export_values();
py::class_<
logging::LockingLogger,
logging::LoggerBase,
std::shared_ptr<logging::LockingLogger>>(m, "LockingLogger")
.def(py::init<>())
.def("set_aggregation_type", &logging::LockingLogger::setAggregationType)
.def("get_counter_val", &logging::LockingLogger::getCounterValue);
py::class_<
logging::NoopLogger,
logging::LoggerBase,
std::shared_ptr<logging::NoopLogger>>(m, "NoopLogger")
.def(py::init<>());
m.def("_jit_is_script_object", [](const py::object& obj) {
return py::isinstance<Object>(obj);
});
initScriptDictBindings(module);
initScriptListBindings(module);
}
} // namespace jit
} // namespace torch