Merge pull request #320 from ROCmSoftwarePlatform/gather_op

Gather op

src/tf/tf.cpp

@@ -174,6 +174,7 @@ struct tf_parser
         add_mem_op("DepthwiseConv2dNative", &tf_parser::parse_depthwiseconv);
         add_mem_op("ExpandDims", &tf_parser::parse_expanddims, false);
         add_mem_op("FusedBatchNorm", &tf_parser::parse_batchnorm);
+        add_mem_op("GatherV2", &tf_parser::parse_gather, false);
         add_mem_op("MatMul", &tf_parser::parse_matmul, false);
         add_mem_op("MaxPool", &tf_parser::parse_pooling);
         add_mem_op("Mean", &tf_parser::parse_mean);
@@ -525,6 +526,14 @@ struct tf_parser
         return prog.add_instruction(op::reshape{new_dims}, args[0]);
     }
 
+    instruction_ref
+    parse_gather(const std::string&, const attribute_map&, std::vector<instruction_ref> args)
+    {
+        int axis = args[2]->eval().at<int32_t>();
+        op::gather op{axis};
+        return prog.add_instruction(op, {args[0], args[1]});
+    }
+
     instruction_ref
     parse_matmul(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
     {

test/tf/tf_test.cpp

@@ -209,6 +209,22 @@ TEST_CASE(expanddims_test_neg_dims)
     EXPECT(p == prog);
 }
 
+TEST_CASE(gather_test)
+{
+    migraphx::program p;
+
+    auto l0 = p.add_parameter("0", migraphx::shape{migraphx::shape::float_type, {2, 4}});
+    auto l1 =
+        p.add_literal(migraphx::literal{migraphx::shape{migraphx::shape::int32_type, {2}}, {1, 1}});
+    p.add_literal(1);
+
+    int axis = 1;
+    p.add_instruction(migraphx::op::gather{axis}, l0, l1);
+    auto prog = optimize_tf("gather_test.pb", false);
+
+    EXPECT(p == prog);
+}
+
 TEST_CASE(identity_test)
 {
     migraphx::program p;