pytorch/torch/csrc/jit/import_source.h
Michael Suo ee9c8a75f4 refactor self to be a class again (#22207)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/22207
ghimport-source-id: 36ee8bd17411a2e220665ad2a27364653061070e

Test Plan: Imported from OSS

Differential Revision: D15998758

Pulled By: suo

fbshipit-source-id: 14bad87bb6e44bf1a43ae86339d8cc7b311c76dd
2019-07-10 15:19:07 -07:00

57 lines
2.2 KiB
C++

#pragma once
#include <torch/csrc/jit/ir.h>
#include <torch/csrc/jit/script/module.h>
#include <functional>
#include <memory>
#include <string>
#include <vector>
namespace torch {
namespace jit {
namespace script {
// Helpers to define modules and classes from source strings. Used in model
// saving/loading, so it expects the format produced by the model exporter.
// Add the methods defined in `src` to the module `mod`.
TORCH_API void import_methods(
// CompilationUnit in which to look up any classes used
const CompilationUnit& lib_cu,
const script::Module& mod,
const std::shared_ptr<Source>& src,
const std::vector<at::Tensor>& constant_table,
// Callback to import any dependencies of this source before compiling
const std::function<void(const std::string&)>& import_callback);
// Define the list of classes in `src`.
TORCH_API void import_libs(
// The compilation unit that will own the imported libs
CompilationUnit& lib_cu,
// Qualifier for any classes that `src` defines. Looks like a module path,
// like "foo.bar.baz"
const std::string& class_qualifier,
const std::shared_ptr<Source>& src,
const std::vector<at::Tensor>& constant_table,
// Callback to import any dependencies of this source before compiling
const std::function<void(const std::string&)>& import_callback);
// Add the functions defined in `src` to the compilation unit `cu`.
// self is passed through the CompilationUnit's define function.
// If present, it determines the SugaredValue for the first argument
// and that argument is no longer expected to have type annotations.
TORCH_API void import_functions(
// Prefix to use when importing these functions in to the CU
const c10::optional<c10::QualifiedName>& prefix,
// CompilationUnit in which to look up any classes used
const CompilationUnit& lib_cu,
// CompilationoUnit to define the functions in.
CompilationUnit& cu,
const std::shared_ptr<Source>& src,
const std::vector<at::Tensor>& constant_table,
const Self* self = nullptr,
const std::function<void(const std::string&)>& import_callback = nullptr);
} // namespace script
} // namespace jit
} // namespace torch