Files
pytorch/test/cpp/jit/test_save_load.cpp
Michael Suo dfdb86a595 big cpp test reorg (#24801)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/24801

This is to fix the ODR-violations in fbcode static builds, which have been broken for several months.

This PR is unfortunately quite large, but the changes are only mechanical:
1. Tests defined in header files -> tests defined in cpp files
2. Remove the `torch::jit::testing` namespace -> `torch::jit`.
3. Single `test.h` file that aggregates all tests.
4. Separate out files for gtest and python versions of the tests instead of using a build flag
5. Add a readme for how to add a new test, and explaining a bit about why the cpp tests are the way they are.

Test Plan: Imported from OSS

Differential Revision: D16878605

Pulled By: suo

fbshipit-source-id: 27b5c077dadd990a5f74e25d01731f9c1f491603
2019-08-18 16:49:56 -07:00

77 lines
1.8 KiB
C++

#include <test/cpp/jit/test_base.h>
#include <test/cpp/jit/test_utils.h>
#include <sstream>
#include <torch/csrc/jit/export.h>
#include <torch/csrc/jit/import.h>
#include <torch/csrc/jit/import_source.h>
#include <torch/torch.h>
namespace torch {
namespace jit {
using namespace script;
void testSaveExtraFilesHook() {
// no secrets
{
std::stringstream ss;
{
Module m("__torch__.m");
ExtraFilesMap extra;
extra["metadata.json"] = "abc";
m.save(ss, extra);
}
ss.seekg(0);
{
ExtraFilesMap extra;
extra["metadata.json"] = "";
extra["secret.json"] = "";
jit::load(ss, c10::nullopt, extra);
ASSERT_EQ(extra["metadata.json"], "abc");
ASSERT_EQ(extra["secret.json"], "");
}
}
// some secret
{
std::stringstream ss;
{
SetExportModuleExtraFilesHook([](const Module&) -> ExtraFilesMap {
return {{"secret.json", "topsecret"}};
});
Module m("__torch__.m");
ExtraFilesMap extra;
extra["metadata.json"] = "abc";
m.save(ss, extra);
SetExportModuleExtraFilesHook(nullptr);
}
ss.seekg(0);
{
ExtraFilesMap extra;
extra["metadata.json"] = "";
extra["secret.json"] = "";
jit::load(ss, c10::nullopt, extra);
ASSERT_EQ(extra["metadata.json"], "abc");
ASSERT_EQ(extra["secret.json"], "topsecret");
}
}
}
static const auto pretty_printed = R"JIT(
op_version_set = 1000
def foo(x: Tensor,
y: Tensor) -> Tensor:
_0 = torch.add(torch.mul(x, 2), y, alpha=1)
return _0
)JIT";
void testImportTooNew() {
Module m("__torch__.m");
const std::vector<at::Tensor> constant_table;
auto src = std::make_shared<Source>(pretty_printed);
ASSERT_ANY_THROW(LEGACY_import_methods(m, src, constant_table, nullptr));
}
} // namespace jit
} // namespace torch