pytorch/test/cpp/jit/test_schema_matching.cpp
Michael Suo 22401b850b port all JIT tests to gtest (#45264)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/45264

Context for why we are porting to gtest in: https://github.com/pytorch/pytorch/pull/45018.

This PR completes the process of porting and removes unused files/macros.

Test Plan: Imported from OSS

Reviewed By: ZolotukhinM

Differential Revision: D23901392

Pulled By: suo

fbshipit-source-id: 89526890e1a49462f3f77718f4ee273c5bc578ba
2020-09-25 11:37:43 -07:00

90 lines
2.1 KiB
C++

#include <gtest/gtest.h>
#include <torch/csrc/jit/ir/ir.h>
#include <torch/csrc/jit/runtime/custom_operator.h>
#include <torch/csrc/jit/testing/file_check.h>
#include <torch/jit.h>
#include <sstream>
#include <string>
namespace torch {
namespace jit {
TEST(SchemaMatchingTest, VarType) {
RegisterOperators reg({
Operator(
"aten::test_vartype(t[] a, t b) -> (t)",
[](Stack* stack) {
c10::List<double> list;
double a;
pop(stack, list, a);
push(stack, a);
},
c10::AliasAnalysisKind::FROM_SCHEMA),
});
Module m("m");
m.define(R"(
def test(self):
a = (1.0, 2.0)
return torch.test_vartype(a, 2.0)
)");
auto result = m.run_method("test");
TORCH_INTERNAL_ASSERT(result.toDouble() == 2.0);
const std::string error_example = R"JIT(
def test_2(self):
a = (1.0, 2.0)
non_float = (1, 1)
return torch.test_vartype(a, non_float)
)JIT";
std::string err = "";
try {
m.define(error_example);
} catch (const std::exception& e) {
err = e.what();
}
TORCH_INTERNAL_ASSERT(
err.find("previously matched to type") != std::string::npos);
}
TEST(SchemaMatchingTest, VarType2) {
RegisterOperators reg({
Operator(
"aten::test_vartype2(t a, t[] b) -> (t[])",
[](Stack* stack) {
double a;
c10::List<double> list;
pop(stack, a, list);
push(stack, a);
},
AliasAnalysisKind::FROM_SCHEMA),
});
Module m("m");
m.define(R"JIT(
def test(self):
a = (1.0, 2.0)
return torch.test_vartype2(3.0, a)
)JIT");
auto result = m.run_method("test");
TORCH_INTERNAL_ASSERT(result.toDouble() == 3.0);
static const auto error_exam2 = R"JIT(
def test_2(self):
a = (1, 2)
return torch.test_vartype2(3.0, a)
)JIT";
std::string err = "";
try {
m.define(error_exam2);
} catch (const std::exception& e) {
err = e.what();
}
TORCH_INTERNAL_ASSERT(
err.find("previously matched to type") != std::string::npos);
}
} // namespace jit
} // namespace torch