add more tests and fixes to parse_axes

src/tf/tf.cpp

@@ -53,15 +53,16 @@ struct tf_parser
     template <class T>
     std::vector<T> parse_axes(std::vector<T> axes) const
     {
-        std::vector<T> new_axes;
         if(is_nhwc)
         {
+            std::vector<T> new_axes;
             std::transform(axes.begin(),
                            axes.end(),
                            std::back_inserter(new_axes),
                            [&](size_t axis) { return parse_axis(axis); });
+            return new_axes;
         }
-        return new_axes;
+        return axes;
     }
 
     // tf stores certain attributes such as strides, dilations, as a 4D input.
@@ -426,7 +427,6 @@ struct tf_parser
     instruction_ref
     parse_mean(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
     {
-
         auto axes      = parse_axes(args[1]->eval().get<int32_t>().to_vector());
         bool keep_dims = attributes.at("keep_dims").b();
         std::vector<int32_t> hw_axes{2, 3};

test/tf/tf_test.cpp

@@ -168,6 +168,40 @@ TEST_CASE(matmul_test)
     EXPECT(p == prog);
 }
 
+TEST_CASE(mean_test)
+{
+    migraphx::program p;
+    migraphx::literal l{migraphx::shape{migraphx::shape::int32_type, {2}}, {2,3}};
+    auto l0 = p.add_parameter("0", migraphx::shape{migraphx::shape::float_type, {1, 3, 16, 16}});
+    p.add_literal(l);
+    p.add_literal(l);
+    migraphx::op::pooling op;
+    op.lengths = {16,16};
+    auto l3 = p.add_instruction(op, l0);
+    p.add_instruction(migraphx::op::squeeze{{2,3}}, l3);
+    p.add_instruction(op, l0);
+    auto prog = migraphx::parse_tf("mean_test.pb", false);
+
+    EXPECT(p == prog);
+}
+
+TEST_CASE(mean_test_nhwc)
+{
+    migraphx::program p;
+    migraphx::literal l{migraphx::shape{migraphx::shape::int32_type, {2}}, {1,2}};
+    auto l0 = p.add_parameter("0", migraphx::shape{migraphx::shape::float_type, {1, 3, 16, 16}});
+    p.add_literal(l);
+    p.add_literal(l);
+    migraphx::op::pooling op;
+    op.lengths = {16,16};
+    auto l3 = p.add_instruction(op, l0);
+    p.add_instruction(migraphx::op::squeeze{{2,3}}, l3);
+    p.add_instruction(op, l0);
+    auto prog = migraphx::parse_tf("mean_test_nhwc.pb", true);
+
+    EXPECT(p == prog);
+}
+
 TEST_CASE(mul_test)
 {
     migraphx::program p;