Commit 37849092 authored by Amy Zhuang's avatar Amy Zhuang

Address PR feedback.

parent bb400712
......@@ -105,7 +105,7 @@ void pass::ConstantFolding::construct_constant_gather()
auto indices_label =
make_shared<pattern::op::Label>(element::i64, Shape{5}, pattern::has_class<op::Constant>());
size_t gather_axis = 1;
auto gather_v0 = make_shared<op::Gather>(data_label, indices_label, gather_axis);
auto gather_v0 = make_shared<op::v0::Gather>(data_label, indices_label, gather_axis);
auto axis_label =
make_shared<pattern::op::Label>(element::i64, Shape{1}, pattern::has_class<op::Constant>());
......
......@@ -1399,14 +1399,14 @@ TEST(constant_folding, const_gather)
auto constant_indices =
op::Constant::create(element::i64, Shape{4}, vector<int64_t>{0, 3, 2, 2});
size_t gather_axis = 1;
auto gather = make_shared<op::Gather>(constant_data, constant_indices, gather_axis);
auto gather = make_shared<op::v0::Gather>(constant_data, constant_indices, gather_axis);
auto f = make_shared<Function>(gather, ParameterVector{});
pass::Manager pass_manager;
pass_manager.register_pass<pass::ConstantFolding>();
pass_manager.run_passes(f);
ASSERT_EQ(count_ops_of_type<op::Gather>(f), 0);
ASSERT_EQ(count_ops_of_type<op::v0::Gather>(f), 0);
ASSERT_EQ(count_ops_of_type<op::Constant>(f), 1);
auto new_const = as_type_ptr<op::Constant>(f->get_results().at(0)->get_argument(0));
......@@ -1434,7 +1434,7 @@ TEST(constant_folding, const_gather_v1)
pass_manager.register_pass<pass::ConstantFolding>();
pass_manager.run_passes(f);
ASSERT_EQ(count_ops_of_type<op::Gather>(f), 0);
ASSERT_EQ(count_ops_of_type<op::v1::Gather>(f), 0);
ASSERT_EQ(count_ops_of_type<op::Constant>(f), 1);
auto new_const = as_type_ptr<op::Constant>(f->get_results().at(0)->get_argument(0));
......@@ -1462,7 +1462,7 @@ TEST(constant_folding, const_gather_v1_scalar)
pass_manager.register_pass<pass::ConstantFolding>();
pass_manager.run_passes(f);
ASSERT_EQ(count_ops_of_type<op::Gather>(f), 0);
ASSERT_EQ(count_ops_of_type<op::v1::Gather>(f), 0);
ASSERT_EQ(count_ops_of_type<op::Constant>(f), 1);
auto new_const = as_type_ptr<op::Constant>(f->get_results().at(0)->get_argument(0));
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment