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
|
#include <gtest/gtest.h>
#include <test/cpp/jit/test_utils.h>
#include <torch/csrc/jit/ir/irparser.h>
#include <torch/csrc/jit/passes/add_if_then_else.h>
namespace torch {
namespace jit {
TEST(AddIfThenElseOpTest, AddIfThenElseOpSimple) {
const auto src = R"IR(
graph(%cond: bool, %a: Tensor, %b: Tensor):
%result: Tensor = prim::If(%cond)
block0():
-> (%a)
block1():
-> (%b)
return (%result)
)IR";
auto graph = std::make_shared<Graph>();
parseIR(src, graph.get());
EXPECT_TRUE(AddIfThenElseOp(graph));
testing::FileCheck()
.check_count("= prim::IfThenElse", 1, /*exactly*/ true)
->check_count("= prim::If", 0, /*exactly*/ true)
->run(*graph);
}
TEST(AddIfThenElseOpTest, NoIfThenElseOpMultipleOutputs) {
const auto src = R"IR(
graph(%cond: bool, %a: Tensor, %b: Tensor):
%result1: Tensor, %result2: Tensor = prim::If(%cond)
block0():
-> (%a, %b)
block1():
-> (%b, %a)
return (%result1, %result2)
)IR";
auto graph = std::make_shared<Graph>();
parseIR(src, graph.get());
EXPECT_FALSE(AddIfThenElseOp(graph));
testing::FileCheck()
.check_count("= prim::IfThenElse", 0, /*exactly*/ true)
->check_count("= prim::If", 1, /*exactly*/ true)
->run(*graph);
}
} // namespace jit
} // namespace torch
|
