Merge pull request #312 from ROCmSoftwarePlatform/transpose_op

Transpose op

src/tf/tf.cpp

@@ -179,6 +179,7 @@ struct tf_parser
         add_mem_op("Softmax", &tf_parser::parse_softmax);
         add_mem_op("Squeeze", &tf_parser::parse_squeeze, false);
         add_mem_op("StridedSlice", &tf_parser::parse_stridedslice);
+        add_mem_op("Transpose", &tf_parser::parse_transpose, false);
     }
 
     template <class F>
@@ -766,6 +767,16 @@ struct tf_parser
         return to_nhwc(prog.add_instruction(op::squeeze{squeeze_axes}, l0));
     }
 
+    instruction_ref
+    parse_transpose(const std::string&, const attribute_map&, std::vector<instruction_ref> args)
+    {
+        auto perm = args[1]->eval().get<int32_t>().to_vector();
+        op::transpose op;
+        op.dims = std::vector<int64_t>(perm.begin(), perm.end());
+
+        return prog.add_instruction(op, args.front());
+    }
+
     void parse_graph(const tensorflow::GraphDef& graph)
     {
         nodes = get_nodes(graph, input_nodes);

test/tf/tf_test.cpp

@@ -435,4 +435,16 @@ TEST_CASE(tanh_test)
     EXPECT(p == prog);
 }
 
+TEST_CASE(transpose_test)
+{
+    migraphx::program p;
+    auto l0 = p.add_parameter("0", migraphx::shape{migraphx::shape::float_type, {1, 3, 16, 16}});
+    migraphx::shape s0{migraphx::shape::int32_type, {4}};
+    p.add_literal(migraphx::literal{s0, {0, 2, 3, 1}});
+    p.add_instruction(migraphx::op::transpose{{0, 2, 3, 1}}, l0);
+    auto prog = optimize_tf("transpose_test.pb", false);
+
+    EXPECT(p == prog);
+}
+
 int main(int argc, const char* argv[]) { test::run(argc, argv); }