add decorators

src/onnx/onnx.cpp

@@ -206,6 +206,16 @@ struct onnx_parser
         return out_lens;
     }
 
+    instruction_ref make_contiguous(instruction_ref ins)
+    {
+        if(ins->get_shape().standard())
+        {
+            return ins;
+        }
+
+        return prog.add_instruction(op::contiguous{}, ins);
+    }
+
     template <class T>
     instruction_ref add_broadcastable_binary_op(instruction_ref arg0, instruction_ref arg1, T x)
     {
@@ -441,12 +451,7 @@ struct onnx_parser
             s.visit([&](auto v) { copy(v, std::back_inserter(op.dims)); });
         }
 
-        if(!args[0]->get_shape().standard())
-        {
-            args[0] = prog.add_instruction(op::contiguous{}, args[0]);
-        }
-
-        return prog.add_instruction(op, args[0]);
+        return prog.add_instruction(op, make_contiguous(args[0]));
     }
 
     instruction_ref
@@ -494,23 +499,41 @@ struct onnx_parser
         {
             axis = parse_value(attributes.at("axis")).at<int>();
         }
+
         op::gather op{axis};
-        return prog.add_instruction(op, std::move(args));
+        return prog.add_instruction(op, make_contiguous(args[0]), make_contiguous(args[1]));
     }
 
     instruction_ref
     parse_slice(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
     {
         op::slice op;
+        std::vector<size_t> dims = args[0]->get_shape().lens();
+        size_t num_dims          = dims.size();
         if(contains(attributes, "axes"))
         {
             literal s = parse_value(attributes.at("axes"));
             s.visit([&](auto v) { copy(v, std::back_inserter(op.axes)); });
         }
+        else
+        {
+            op.axes = std::vector<int64_t>(num_dims);
+            std::iota(op.axes.begin(), op.axes.end(), 0);
+        }
+
+        if(contains(attributes, "ends"))
         {
             literal s = parse_value(attributes.at("ends"));
             s.visit([&](auto v) { copy(v, std::back_inserter(op.ends)); });
+            for(size_t i = 0; i < num_dims; i++)
+            {
+                if(static_cast<size_t>(op.ends[i]) > dims[i])
+                {
+                    op.ends[i] = dims[i];
+                }
+            }
         }
+        if(contains(attributes, "starts"))
         {
             literal s = parse_value(attributes.at("starts"));
             s.visit([&](auto v) { copy(v, std::back_inserter(op.starts)); });

src/quantization.cpp

@@ -74,7 +74,8 @@ void quantize(program& prog, const std::vector<std::string>& ins_names)
                 // if the input is a convert operator, uses its input
                 // as its current input
                 instruction_ref input_fp16{};
-                if(input->name() == "convert")
+                if(input->name() == "convert" and
+                   input->inputs().front()->get_shape().type() == shape::half_type)
                 {
                     input_fp16 = input->inputs().front();
                 }

src/targets/gpu/quant_gemm.cpp

@@ -8,51 +8,6 @@ namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
 namespace gpu {
 
-template <class... Ts>
-rocblas_status generic_rocblas_gemm_ex(Ts&&... xs)
-{
-    return rocblas_gemm_ex(std::forward<Ts>(xs)...);
-}
-
-template <class... Ts>
-rocblas_status generic_rocblas_batched_gemm_ex(Ts&&... xs)
-{
-    return rocblas_gemm_strided_batched_ex(std::forward<Ts>(xs)...);
-}
-
-template <class T>
-struct compute_rocblas_type
-{
-    using type = T;
-};
-
-template <class T>
-struct compute_rocblas_type<const T>
-{
-    using type = const typename compute_rocblas_type<T>::type;
-};
-
-template <>
-struct compute_rocblas_type<half>
-{
-    using type = rocblas_half;
-};
-
-template <class T>
-using rb_type = typename compute_rocblas_type<T>::type;
-
-template <class T>
-rb_type<T> to_rocblas_type(T x)
-{
-    return reinterpret_cast<const rb_type<T>&>(x);
-}
-
-template <class T>
-rb_type<T>* to_rocblas_type(T* x)
-{
-    return reinterpret_cast<rb_type<T>*>(x);
-}
-
 shape rocblas_quant_gemm::compute_shape(const std::vector<shape>& inputs) const
 {
     std::vector<shape> in_shapes(inputs);
@@ -102,13 +57,13 @@ argument rocblas_quant_gemm::compute(context& ctx,
     auto a_lens = args[0].get_shape().lens();
     auto b_lens = args[1].get_shape().lens();
     output_shape.visit_type([&](auto as) {
-        auto alpha_r    = to_rocblas_type(as(op.alpha));
-        auto beta_r     = to_rocblas_type(as(beta));
+        auto alpha_r    = as(op.alpha);
+        auto beta_r     = as(beta);
         auto out_lens   = output_shape.lens();
         rocblas_int m   = out_lens[dim_0];
         rocblas_int n   = out_lens[dim_1];
         rocblas_int k   = args[0].get_shape().lens()[dim_1];
-        auto to_pointer = [&](auto&& arg) { return to_rocblas_type(as.from(arg.data())); };
+        auto to_pointer = [&](auto&& arg) { return as.from(arg.data()); };
         assert(k % 4 == 0);
 
         auto num_matrices = std::accumulate(
@@ -119,36 +74,36 @@ argument rocblas_quant_gemm::compute(context& ctx,
             // column-major format. When doing a C = A * B, we actually do
             // C^T = (B^T) * (A^T). That is the reason we input args[1] as
             // A and args[0] as B in calling the rocblas_gemm.
-            generic_rocblas_gemm_ex(ctx.get_stream().get_rocblas(),
-                                    transb ? rocblas_operation_transpose : rocblas_operation_none,
-                                    transa ? rocblas_operation_transpose : rocblas_operation_none,
-                                    n,
-                                    m,
-                                    k,
-                                    &alpha_r,
-                                    to_pointer(args.at(1)),
-                                    rocblas_datatype_i8_r,
-                                    ldb,
-                                    to_pointer(args.at(0)),
-                                    rocblas_datatype_i8_r,
-                                    lda,
-                                    &beta_r,
-                                    to_pointer(args[2]),
-                                    rocblas_datatype_i32_r,
-                                    ldc,
-                                    is_3inputs ? to_pointer(args[3]) : to_pointer(args[2]),
-                                    rocblas_datatype_i32_r,
-                                    ldc,
-                                    rocblas_datatype_i32_r,
-                                    rocblas_gemm_algo_standard,
-                                    0,
-                                    0,
-                                    nullptr,
-                                    nullptr);
+            rocblas_gemm_ex(ctx.get_stream().get_rocblas(),
+                            transb ? rocblas_operation_transpose : rocblas_operation_none,
+                            transa ? rocblas_operation_transpose : rocblas_operation_none,
+                            n,
+                            m,
+                            k,
+                            &alpha_r,
+                            to_pointer(args.at(1)),
+                            rocblas_datatype_i8_r,
+                            ldb,
+                            to_pointer(args.at(0)),
+                            rocblas_datatype_i8_r,
+                            lda,
+                            &beta_r,
+                            to_pointer(args[2]),
+                            rocblas_datatype_i32_r,
+                            ldc,
+                            is_3inputs ? to_pointer(args[3]) : to_pointer(args[2]),
+                            rocblas_datatype_i32_r,
+                            ldc,
+                            rocblas_datatype_i32_r,
+                            rocblas_gemm_algo_standard,
+                            0,
+                            0,
+                            nullptr,
+                            nullptr);
         }
         else
         {
-            generic_rocblas_batched_gemm_ex(
+            rocblas_gemm_strided_batched_ex(
                 ctx.get_stream().get_rocblas(),
                 transb ? rocblas_operation_transpose : rocblas_operation_none,
                 transa ? rocblas_operation_transpose : rocblas_operation_none,

src/tf/tf.cpp

@@ -26,7 +26,6 @@ struct tf_parser
 {
     using attribute_map = std::unordered_map<std::string, tensorflow::AttrValue>;
     using node_map      = std::map<std::string, tensorflow::NodeDef>;
-    // using input_node_map = std::unordered_map<std::string, std::unordered_set<std::string>>;
     using op_func = std::function<instruction_ref(attribute_map, std::vector<instruction_ref>)>;
 
     node_map nodes;
@@ -149,9 +148,26 @@ struct tf_parser
         return axes;
     }
 
+    std::vector<int64_t> get_axes_from_mask(const size_t num_axes, const uint32_t mask)
+    {
+        uint32_t bitwise_compare = 1;
+        std::vector<int64_t> axes;
+        for(size_t i = 0; i < num_axes; i++)
+        {
+            // the LSB corresponds to axis 0 when determining which axes to begin
+            if(((mask >> i) & bitwise_compare) == 1)
+                axes.push_back(1);
+            else
+                axes.push_back(0);
+        }
+        return axes;
+    }
+
     tf_parser()
     {
+        add_generic_op("All", op::identity{});
         add_generic_op("Identity", op::identity{});
+        add_generic_op("LessEqual", op::identity{});
         add_generic_op("Relu", op::relu{});
         add_generic_op("Relu6", op::clip{6.0, 0.0});
         add_generic_op("Rsqrt", op::rsqrt{});
@@ -166,6 +182,7 @@ struct tf_parser
 
         add_mem_op("AvgPool", &tf_parser::parse_pooling);
         add_mem_op("BatchMatMul", &tf_parser::parse_matmul, false);
+        add_mem_op("BatchMatMulV2", &tf_parser::parse_matmul, false);
         add_mem_op("BiasAdd", &tf_parser::parse_biasadd);
         add_mem_op("Cast", &tf_parser::parse_cast, false);
         add_mem_op("ConcatV2", &tf_parser::parse_concat, false);
@@ -177,14 +194,15 @@ struct tf_parser
         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);
+        add_mem_op("Mean", &tf_parser::parse_mean, false);
+        add_mem_op("OneHot", &tf_parser::parse_onehot, false);
         add_mem_op("Pack", &tf_parser::parse_pack, false);
         add_mem_op("Pad", &tf_parser::parse_pad);
         add_mem_op("Reshape", &tf_parser::parse_reshape, false);
         add_mem_op("Slice", &tf_parser::parse_slice, false);
-        add_mem_op("Softmax", &tf_parser::parse_softmax<op::softmax>);
+        add_mem_op("Softmax", &tf_parser::parse_softmax<op::softmax>, false);
         add_mem_op("Squeeze", &tf_parser::parse_squeeze, false);
-        add_mem_op("StridedSlice", &tf_parser::parse_stridedslice);
+        add_mem_op("StridedSlice", &tf_parser::parse_stridedslice, false);
         add_mem_op("Transpose", &tf_parser::parse_transpose, false);
     }
 
@@ -547,7 +565,7 @@ struct tf_parser
         }
         if(contains(attributes, "transpose_b"))
         {
-            transb = attributes.at("transpose_a").b();
+            transb = attributes.at("transpose_b").b();
         }
 
         if(contains(attributes, "adj_x"))
@@ -574,8 +592,7 @@ struct tf_parser
     parse_mean(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
     {
         bool keep_dims = attributes.at("keep_dims").b();
-        auto lens      = args[0]->get_shape().lens();
-        auto axes = parse_axes(args[1]->eval().get<int32_t>().to_vector<int64_t>(), lens.size());
+        auto axes      = args[1]->eval().get<int32_t>().to_vector<int64_t>();
 
         if(keep_dims)
         {
@@ -588,6 +605,32 @@ struct tf_parser
         }
     }
 
+    instruction_ref
+    parse_onehot(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
+    {
+        size_t depth = static_cast<size_t>(args[1]->eval().at<int32_t>());
+
+        int64_t axis    = -1;
+        float on_value  = args[2]->eval().at<float>();
+        float off_value = args[3]->eval().at<float>();
+
+        std::vector<float> depth_input(depth * depth, off_value);
+        for(int i = 0; i < depth; i++)
+        {
+            depth_input[depth * i + i] = on_value;
+        }
+
+        if(contains(attributes, "axis"))
+            axis = attributes.at("axis").i();
+        if(axis == -1)
+        {
+            shape s{shape::float_type, {depth, depth}};
+            auto l0 = prog.add_literal({s, depth_input});
+            return prog.add_instruction(op::gather{0}, {l0, args[0]});
+        }
+        MIGRAPHX_THROW("MIGraphX does not support axis != -1");
+    }
+
     instruction_ref parse_pack(const std::string&,
                                const attribute_map& attributes,
                                std::vector<instruction_ref> args)
@@ -799,21 +842,50 @@ struct tf_parser
                                        std::vector<instruction_ref> args)
     {
         op::slice op;
-        auto starts     = args[1]->eval().get<int32_t>().to_vector();
-        auto ends       = args[2]->eval().get<int32_t>().to_vector();
-        size_t num_axes = args[0]->get_shape().lens().size();
+        auto starts              = args[1]->eval().get<int32_t>().to_vector();
+        auto ends                = args[2]->eval().get<int32_t>().to_vector();
+        auto l0                  = args[0];
+        size_t num_axes          = l0->get_shape().lens().size();
+        std::vector<size_t> axes = l0->get_shape().lens();
 
         op.starts = std::vector<int64_t>(starts.begin(), starts.end());
         op.ends   = std::vector<int64_t>(ends.begin(), ends.end());
         op.axes   = std::vector<int64_t>(num_axes);
         std::iota(op.axes.begin(), op.axes.end(), 0);
+        uint32_t begin_mask       = 0;
+        uint32_t end_mask         = 0;
         uint32_t shrink_axis_mask = 0;
         uint32_t bitwise_compare  = 1;
         std::vector<int64_t> squeeze_axes;
 
+        if(contains(attributes, "begin_mask"))
+            begin_mask = static_cast<uint32_t>(attributes.at("begin_mask").i());
+
+        if(contains(attributes, "end_mask"))
+            end_mask = static_cast<uint32_t>(attributes.at("end_mask").i());
+
         if(contains(attributes, "shrink_axis_mask"))
             shrink_axis_mask = static_cast<uint32_t>(attributes.at("shrink_axis_mask").i());
 
+        std::vector<int64_t> begin_axes = get_axes_from_mask(num_axes, begin_mask);
+        std::vector<int64_t> end_axes   = get_axes_from_mask(num_axes, end_mask);
+
+        for(size_t i = 0; i < num_axes; i++)
+        {
+            if(begin_axes.at(i) == 1)
+            {
+                op.starts.at(i) = 0;
+            }
+            if(end_axes.at(i) == 1)
+            {
+                op.ends.at(i) = axes.at(i);
+            }
+        }
+
+        auto l1 = prog.add_instruction(op, l0);
+        if(shrink_axis_mask == 0)
+            return l1;
+
         for(size_t i = 0; i < num_axes; i++)
         {
             // the LSB corresponds to axis 0 when determining which axes to squeeze
@@ -821,8 +893,7 @@ struct tf_parser
                 squeeze_axes.push_back(i);
         }
 
-        auto l0 = prog.add_instruction(op, make_contiguous(args[0]));
-        return to_nhwc(prog.add_instruction(op::squeeze{squeeze_axes}, l0));
+        return prog.add_instruction(op::squeeze{squeeze_axes}, l1);
     }
 
     instruction_ref
@@ -862,10 +933,16 @@ struct tf_parser
         if(instructions.count(name) == 0)
         {
             auto&& node = nodes.at(name);
+            // assert ops ignored
+            if(node.op() == "Assert" or contains(name, "Assert"))
+                return;
             std::vector<instruction_ref> args;
 
             for(auto&& input : node.input())
             {
+                // control dependencies (signified by ^ before the name) are ignored
+                if(contains(input, "^"))
+                    continue;
                 if(nodes.count(input) > 0)
                 {
                     auto&& iname = get_name(nodes.at(input));

test/onnx/onnx_test.cpp

@@ -4,6 +4,7 @@
 #include <migraphx/operators.hpp>
 #include <migraphx/program.hpp>
 #include <migraphx/instruction.hpp>
+#include <migraphx/instruction_ref.hpp>
 #include <migraphx/onnx.hpp>
 #include "test.hpp"
 
@@ -1015,6 +1016,32 @@ TEST_CASE(transpose_test)
     EXPECT(p == prog);
 }
 
+TEST_CASE(transpose_gather_test)
+{
+    migraphx::program p;
+    auto make_contiguous = [&p](migraphx::instruction_ref ins) {
+        if(ins->get_shape().standard())
+        {
+            return ins;
+        }
+
+        return p.add_instruction(migraphx::op::contiguous{}, ins);
+    };
+
+    auto data = p.add_parameter("data", migraphx::shape{migraphx::shape::float_type, {3, 5, 4, 6}});
+    auto ind =
+        p.add_parameter("indices", migraphx::shape{migraphx::shape::int32_type, {2, 4, 3, 5}});
+    auto tr_data = p.add_instruction(migraphx::op::transpose{{0, 2, 1, 3}}, data);
+    auto tr_ind  = p.add_instruction(migraphx::op::transpose{{0, 2, 1, 3}}, ind);
+    int axis     = 1;
+    p.add_instruction(
+        migraphx::op::gather{axis}, make_contiguous(tr_data), make_contiguous(tr_ind));
+
+    auto prog = migraphx::parse_onnx("transpose_gather_test.onnx");
+
+    EXPECT(p == prog);
+}
+
 TEST_CASE(unknown_test)
 {
     migraphx::program p;

test/tf/gen_tf_pb.py

 import numpy as np
 import tensorflow as tf
 
+def tf_test(op_test):
+    def run_test():
+        g1 = tf.Graph()
+        op_test(g1)
+        tf.io.write_graph(g1, '.', '{}.pb'.format(op_test.__name__), as_text=False)
+    return run_test
 
-def add_test(g1=tf.Graph()):
+@tf_test
+def add_test(g1):
     with g1.as_default():
         g1_input = tf.placeholder(tf.float32, shape=(1,2,2,3), name = '0')
         g2_input = tf.placeholder(tf.float32, shape=(1,2,2,3), name = '1')
         tf.add(g1_input, g2_input, name = 'add1')
 
-    tf.train.write_graph(g1, '.', 'add_test.pb', as_text=False)
-
-def add_bcast_test(g1=tf.Graph()):
+@tf_test
+def add_bcast_test(g1):
     with g1.as_default():
         g1_input = tf.placeholder(tf.float32, shape=(2,3), name = '0')
         g2_input = tf.placeholder(tf.float32, shape=(2,1), name = '1')
         tf.math.add(g1_input, g2_input, name = 'add_bcast1')
 
-    tf.train.write_graph(g1, '.', 'add_bcast_test.pb', as_text=False)
-
-def assert_less_equal_test(g1=tf.Graph()):
+@tf_test
+def assert_less_equal_test(g1):
     with g1.as_default():
         g1_input = tf.placeholder(tf.float32, shape=(2,3), name = '0')
         g2_input = tf.placeholder(tf.float32, shape=(2,3), name = '1')
         with tf.control_dependencies([tf.assert_less_equal(g1_input, g2_input)]):
             tf.add(g1_input, g2_input, name = 'add1')
-    
-    tf.train.write_graph(g1, '.', 'assert_less_equal_test.pb', as_text=False)
 
-def batchmatmul_test(g1=tf.Graph()):
+@tf_test
+def batchmatmul_test(g1):
     with g1.as_default():
         g1_input = tf.placeholder(tf.float32, shape=(1,2,8,4), name = '0')
         g2_input = tf.placeholder(tf.float32, shape=(1,2,4,8), name = '1')
         tf.matmul(g1_input, g2_input, transpose_a=True, transpose_b=True, name='batchmatmul1')
 
-    tf.train.write_graph(g1, '.', 'batchmatmul_test.pb', as_text=False)
-
-def batchnorm_test(g1=tf.Graph()):
+@tf_test
+def batchnorm_test(g1):
     with g1.as_default():
         g1_input = tf.placeholder(tf.float32, shape=(1, 16, 16, 32), name = '0')
         g1_scale = tf.constant(1.0, dtype=tf.float32, shape=[32], name = '1')
@@ -43,271 +46,241 @@ def batchnorm_test(g1=tf.Graph()):
         g1_mean = tf.placeholder(tf.float32, shape=(32), name = '3')
         g1_variance = tf.placeholder(tf.float32, shape=(32), name = '4')
         tf.nn.fused_batch_norm(
-                g1_input, g1_scale, g1_offset, g1_mean, g1_variance,
-                epsilon=0.00001, is_training=False, name='batchnorm1')
+            g1_input, g1_scale, g1_offset, g1_mean, g1_variance,
+            epsilon=0.00001, is_training=False, name='batchnorm1'
+        )
 
-    tf.train.write_graph(g1, '.', 'batchnorm_test.pb', as_text=False)
-
-def biasadd_test(g1=tf.Graph()):
+@tf_test
+def biasadd_test(g1):
     with g1.as_default():
         g1_input = tf.placeholder(tf.float32, shape=(1,1,1,500), name = '0')
         g2_input = tf.placeholder(tf.float32, shape=(500), name = '1')
         tf.nn.bias_add(g1_input, g2_input, name = 'bias_add1')
 
-    tf.train.write_graph(g1, '.', 'biasadd_test.pb', as_text=False)
-
-def cast_test(g1=tf.Graph()):
+@tf_test
+def cast_test(g1):
     with g1.as_default():
         g1_input = tf.placeholder(tf.float32, shape=(1,3,16,16), name = '0')
         tf.cast(g1_input, dtype=tf.int32, name='cast1')
 
-    tf.train.write_graph(g1, '.', 'cast_test.pb', as_text=False)
-def concat_test(g1=tf.Graph()):
+@tf_test
+def concat_test(g1):
     with g1.as_default():
         g1_input = tf.placeholder(tf.float32, shape=(4,7,3), name = '0')
         g2_input = tf.placeholder(tf.float32, shape=(4,2,3), name = '1')
         tf.concat([g1_input, g2_input], axis=1, name = 'concat1')
 
-    tf.train.write_graph(g1, '.', 'concat_test.pb', as_text=False)
-
-def const_test(g1=tf.Graph()):
+@tf_test
+def const_test(g1):
     with g1.as_default():
         tf.constant(1.0, dtype=tf.float32 ,name='constant1')
 
-    tf.train.write_graph(g1, '.', 'constant_test.pb', as_text=False)
-
-def conv_test(g1=tf.Graph()):
+@tf_test
+def conv_test(g1):
     with g1.as_default():
         g1_input = tf.placeholder(tf.float32, shape=(1,16,16,3), name = '0')
         g1_weights = tf.constant(value=1.0, dtype=tf.float32, shape=(3,3,3,32), name = '1')
         tf.nn.conv2d(g1_input, g1_weights, [1,1,1,1], "SAME", name = 'conv1')
 
-    tf.train.write_graph(g1, '.', 'conv_test.pb', as_text=False)
-
-def depthwiseconv_test(g1=tf.Graph()):
+@tf_test
+def depthwiseconv_test(g1):
     with g1.as_default():
         g1_input = tf.placeholder(tf.float32, shape=(1,16,16,3), name = '0')
         g1_weights = tf.constant(value=1.0, dtype=tf.float32, shape=(3,3,3,1), name = '1')
         tf.nn.depthwise_conv2d_native(g1_input, g1_weights, [1,1,1,1], "SAME", name = 'depthwiseconv1')
 
-    tf.train.write_graph(g1, '.', 'depthwise_conv_test.pb', as_text=False)
-
-def expanddims_test(g1=tf.Graph()):
+@tf_test
+def expanddims_test(g1):
     with g1.as_default():
         g1_input = tf.placeholder(tf.float32, shape=(2,3,4), name = '0')
         tf.expand_dims(g1_input, axis=-1, name='expanddims_neg')
 
-    tf.train.write_graph(g1, '.', 'expanddims_neg_test.pb', as_text=False)
-
-def gather_test(g1=tf.Graph()):
+@tf_test
+def gather_test(g1):
     with g1.as_default():
         g1_input = tf.placeholder(tf.float32, shape=(2,4), name = '0')
         tf.gather(g1_input, [1,1], axis=1, name='gather1')
 
-    tf.train.write_graph(g1, '.', 'gather_test.pb', as_text=False)
-
-def identity_test(g1=tf.Graph()):
+@tf_test
+def identity_test(g1):
     with g1.as_default():
         g1_input = tf.placeholder(tf.float32, shape=(1,3,16,16), name = '0')
         tf.identity(g1_input, 'identity')
 
-    tf.train.write_graph(g1, '.', 'identity_test.pb', as_text=False)
-
-def matmul_test(g1=tf.Graph()):
+@tf_test
+def matmul_test(g1):
     with g1.as_default():
         g1_input = tf.placeholder(tf.float32, shape=(8,4), name = '0')
         g2_input = tf.placeholder(tf.float32, shape=(4,8), name = '1')
         tf.matmul(g1_input, g2_input, transpose_a=True, transpose_b=True, name='matmul1')
 
-    tf.train.write_graph(g1, '.', 'matmul_test.pb', as_text=False)
-
-def mean_test(g1=tf.Graph()):
+@tf_test
+def mean_test(g1):
     with g1.as_default():
         g1_input = tf.placeholder(tf.float32, shape=(1,3,16,16), name = '0')
         tf.math.reduce_mean(
-                g1_input,
-                axis=(2,3),
-                keepdims=True,
-                name='mean1'
-            )
+            g1_input,
+            axis=(2,3),
+            keepdims=True,
+            name='mean1'
+        )
         tf.math.reduce_mean(
-                g1_input,
-                axis=(2,3),
-                keepdims=False,
-                name='mean2'
-            )
+            g1_input,
+            axis=(2,3),
+            keepdims=False,
+            name='mean2'
+        )
 
-    tf.train.write_graph(g1, '.', 'mean_test.pb', as_text=False)
-
-def mean_test_nhwc(g1=tf.Graph()):
+@tf_test
+def mean_test_nhwc(g1):
     with g1.as_default():
         g1_input = tf.placeholder(tf.float32, shape=(1,16,16,3), name = '0')
         tf.math.reduce_mean(
-                g1_input,
-                axis=(1,2),
-                keepdims=True,
-                name='mean1'
-            )
+            g1_input,
+            axis=(1,2),
+            keepdims=True,
+            name='mean1'
+        )
         tf.math.reduce_mean(
-                g1_input,
-                axis=(1,2),
-                keepdims=False,
-                name='mean2'
-            )
-
-    tf.train.write_graph(g1, '.', 'mean_test_nhwc.pb', as_text=False)
+            g1_input,
+            axis=(1,2),
+            keepdims=False,
+            name='mean2'
+        )
 
-def mul_test(g1=tf.Graph()):
+@tf_test
+def mul_test(g1):
     with g1.as_default():
         g1_input = tf.placeholder(tf.float32, shape=(1,1,1,16), name = '0')
         g2_input = tf.placeholder(tf.float32, shape=(1,1,1,16), name = '1')
         tf.multiply(g1_input, g2_input, name='mul1')
-    tf.train.write_graph(g1, '.', 'mul_test.pb', as_text=False)
 
-def pack_test(g1=tf.Graph()):
+@tf_test
+def pack_test(g1):
     with g1.as_default():
         g1_input = tf.placeholder(tf.float32, shape=(2), name = '0')
         g2_input = tf.placeholder(tf.float32, shape=(2), name = '1')
         g3_input = tf.placeholder(tf.float32, shape=(2), name = '2')
         tf.stack([g1_input, g2_input, g3_input], axis=1, name = 'pack1')
 
-    tf.train.write_graph(g1, '.', 'pack_test.pb', as_text=False)
-
-def pack_test_nhwc(g1=tf.Graph()):
+@tf_test
+def pack_test_nhwc(g1):
     with g1.as_default():
         g1_input = tf.placeholder(tf.float32, shape=(1,1,1,2), name = '0')
         g2_input = tf.placeholder(tf.float32, shape=(1,1,1,2), name = '1')
         g3_input = tf.placeholder(tf.float32, shape=(1,1,1,2), name = '2')
         tf.stack([g1_input, g2_input, g3_input], axis=3, name = 'pack1')
 
-    tf.train.write_graph(g1, '.', 'pack_test_nhwc.pb', as_text=False)
-
-def pooling_test(g1=tf.Graph()):
+@tf_test
+def pooling_test(g1):
     with g1.as_default():
         g1_input = tf.placeholder(tf.float32, shape=(1,16,16,3), name = '0')
         tf.nn.avg_pool(
-                value=g1_input,
-                ksize=(1,2,2,1),
-                strides=(1,2,2,1),
-                padding='VALID',
-                data_format='NHWC',
-                name='avg_pooling'
-            )
+            value=g1_input,
+            ksize=(1,2,2,1),
+            strides=(1,2,2,1),
+            padding='VALID',
+            data_format='NHWC',
+            name='avg_pooling'
+        )
         tf.nn.max_pool(
-                value=g1_input,
-                ksize=(1,2,2,1),
-                strides=(1,2,2,1),
-                padding='VALID',
-                data_format='NHWC',
-                name='max_pooling'
-            )
+            value=g1_input,
+            ksize=(1,2,2,1),
+            strides=(1,2,2,1),
+            padding='VALID',
+            data_format='NHWC',
+            name='max_pooling'
+        )
 
-    tf.train.write_graph(g1, '.', 'pooling_test.pb', as_text=False)
-
-def pow_test(g1=tf.Graph()):
+@tf_test
+def pow_test(g1):
     with g1.as_default():
         g1_input = tf.placeholder(tf.float32, shape=(1,2,2,3), name = '0')
         g2_input = tf.placeholder(tf.float32, shape=(1,2,2,3), name = '1')
         tf.pow(g1_input, g2_input, name = 'pow1')
 
-    tf.train.write_graph(g1, '.', 'pow_test.pb', as_text=False)
-
-def relu_test(g1=tf.Graph()):
+@tf_test
+def relu_test(g1):
     with g1.as_default():
         g1_input = tf.placeholder(tf.float32, shape=(1,3,16,16), name = '0')
         tf.nn.relu(g1_input, 'relu')
 
-    tf.train.write_graph(g1, '.', 'relu_test.pb', as_text=False)
-
-def relu6_test(g1=tf.Graph()):
+@tf_test
+def relu6_test(g1):
     with g1.as_default():
         g1_input = tf.placeholder(tf.float32, shape=(1,3,16,16), name = '0')
         tf.nn.relu6(g1_input, 'relu6')
 
-    tf.train.write_graph(g1, '.', 'relu6_test.pb', as_text=False)
-
-def reshape_test(g1=tf.Graph()):
+@tf_test
+def reshape_test(g1):
     with g1.as_default():
         g1_input = tf.placeholder(tf.float32, shape=(16), name = '0')
         tf.reshape(g1_input, (1,1,1,16), 'reshape')
 
-    tf.train.write_graph(g1, '.', 'reshape_test.pb', as_text=False)
-
-def rsqrt_test(g1=tf.Graph()):
+@tf_test
+def rsqrt_test(g1):
     with g1.as_default():
         g1_input = tf.placeholder(tf.float32, shape=(1,3,16,16), name = '0')
         tf.math.rsqrt(g1_input, 'rsqrt')
 
-    tf.train.write_graph(g1, '.', 'rsqrt_test.pb', as_text=False)
-
-def slice_test(g1=tf.Graph()):
+@tf_test
+def slice_test(g1):
     with g1.as_default():
         g1_input = tf.placeholder(tf.float32, shape=(5,10), name = '0')
         tf.slice(g1_input, [1, 0], [2, -1], name = 'slice1')
 
-    tf.train.write_graph(g1, '.', 'slice_test.pb', as_text=False)
-
-def softmax_test(g1=tf.Graph()):
+@tf_test
+def softmax_test(g1):
     with g1.as_default():
         g1_input = tf.placeholder(tf.float32, shape=(1,3), name = '0')
         tf.nn.softmax(g1_input, name='softmax')
 
-    tf.train.write_graph(g1, '.', 'softmax_test.pb', as_text=False)
-
-def sqdiff_test(g1=tf.Graph()):
+@tf_test
+def sqdiff_test(g1):
     with g1.as_default():
         g1_input = tf.placeholder(tf.float32, shape=(1,2,2,3), name = '0')
         g2_input = tf.placeholder(tf.float32, shape=(1,2,2,3), name = '1')
         tf.squared_difference(g1_input, g2_input, name = 'sqdiff')
 
-    tf.train.write_graph(g1, '.', 'sqdiff_test.pb', as_text=False)
-
-def squeeze_test(g1=tf.Graph()):
+@tf_test
+def squeeze_test(g1):
     with g1.as_default():
         g1_input = tf.placeholder(tf.float32, shape=(1,2,3,1), name = '0')
         tf.squeeze(g1_input, name='squeeze')
 
-    tf.train.write_graph(g1, '.', 'squeeze_test.pb', as_text=False)
-
-def stopgradient_test(g1=tf.Graph()):
+@tf_test
+def stopgradient_test(g1):
     with g1.as_default():
         g1_input = tf.placeholder(tf.float32, shape=(1,3,16,16), name = '0')
         tf.stop_gradient(g1_input, 'stopgradient')
 
-    tf.train.write_graph(g1, '.', 'stopgradient_test.pb', as_text=False)
-
-def stridedslice_test(g1=tf.Graph()):
+@tf_test
+def stridedslice_test(g1):
     with g1.as_default():
         g1_input = tf.placeholder(tf.float32, shape=(1,1,1,10), name = '0')
         tf.strided_slice(g1_input, [0, 0, 0, 0], [1, 1, 1, 5], [1,1,1,1], shrink_axis_mask=2, name = 'stridedslice1')
 
-    tf.train.write_graph(g1, '.', 'stridedslice_test.pb', as_text=False)
-
-def stridedslice_masks_test(g1=tf.Graph()):
+@tf_test
+def stridedslice_masks_test(g1):
     with g1.as_default():
         g1_input = tf.placeholder(tf.float32, shape=(1,3,3,10), name = '0')
         tf.strided_slice(g1_input, [0, 1, 1, 0], [0, 0, 0, 0], [1,1,1,1], begin_mask=9, end_mask=15, name = 'stridedslice1')
 
-    tf.train.write_graph(g1, '.', 'stridedslice_masks_test.pb', as_text=False)
-
-def sub_test(g1=tf.Graph()):
+@tf_test
+def sub_test(g1):
     with g1.as_default():
         g1_input = tf.placeholder(tf.float32, shape=(1,2,2,3), name = '0')
         g2_input = tf.placeholder(tf.float32, shape=(1,2,2,3), name = '1')
         tf.subtract(g1_input, g2_input, name = 'sub1')
 
-    tf.train.write_graph(g1, '.', 'sub_test.pb', as_text=False)
-
-def tanh_test(g1=tf.Graph()):
+@tf_test
+def tanh_test(g1):
     with g1.as_default():
         g1_input = tf.placeholder(tf.float32, shape=(1,3,16,16), name = '0')
         tf.tanh(g1_input, 'tanh')
 
-    tf.train.write_graph(g1, '.', 'tanh_test.pb', as_text=False)
-
-def transpose_test(g1=tf.Graph()):
+@tf_test
+def transpose_test(g1):
     with g1.as_default():
         g1_input = tf.placeholder(tf.float32, shape=(1,3,16,16), name = '0')
         tf.transpose(g1_input, perm=[0,2,3,1], name = 'transpose')
-
-    tf.train.write_graph(g1, '.', 'transpose_test.pb', as_text=False)

test/tf/tf_test.cpp

@@ -48,6 +48,22 @@ TEST_CASE(add_bcast_test)
     EXPECT(p == prog);
 }
 
+TEST_CASE(assert_less_equal_test)
+{
+    migraphx::program p;
+    migraphx::shape s0{migraphx::shape::float_type, {2, 3}};
+    auto l0 = p.add_parameter("0", s0);
+    auto l1 = p.add_parameter("1", s0);
+    migraphx::literal l{migraphx::shape{migraphx::shape::int32_type, {2}}, {0, 1}};
+    auto l2 = p.add_literal(l);
+    p.add_instruction(migraphx::op::add{}, l0, l1);
+    auto l3 = p.add_instruction(migraphx::op::identity{}, l0, l1);
+    p.add_instruction(migraphx::op::identity{}, l3, l2);
+    auto prog = optimize_tf("assert_less_equal_test.pb", false);
+
+    EXPECT(p == prog);
+}
+
 TEST_CASE(batchmatmul_test)
 {
     migraphx::program p;
@@ -100,6 +116,16 @@ TEST_CASE(biasadd_test)
     EXPECT(p == prog);
 }
 
+TEST_CASE(cast_test)
+{
+    migraphx::program p;
+    auto l0 = p.add_parameter("0", migraphx::shape{migraphx::shape::float_type, {1, 3, 16, 16}});
+    p.add_instruction(migraphx::op::convert{migraphx::shape::int32_type}, l0);
+    auto prog = optimize_tf("cast_test.pb", false);
+
+    EXPECT(p == prog);
+}
+
 TEST_CASE(concat_test)
 {
     migraphx::program p;
@@ -118,16 +144,6 @@ TEST_CASE(concat_test)
     EXPECT(p == prog);
 }
 
-TEST_CASE(cast_test)
-{
-    migraphx::program p;
-    auto l0 = p.add_parameter("0", migraphx::shape{migraphx::shape::float_type, {1, 3, 16, 16}});
-    p.add_instruction(migraphx::op::convert{migraphx::shape::int32_type}, l0);
-    auto prog = optimize_tf("cast_test.pb", false);
-
-    EXPECT(p == prog);
-}
-
 TEST_CASE(const_test)
 {
     migraphx::program p;
@@ -271,9 +287,10 @@ 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}});
-    migraphx::op::reduce_mean op{{2, 3}};
-    auto l3 = p.add_instruction(op, l0);
-    p.add_instruction(migraphx::op::squeeze{{2, 3}}, l3);
+    auto l1 = p.add_instruction(migraphx::op::transpose{{0, 2, 3, 1}}, l0);
+    migraphx::op::reduce_mean op{{1, 2}};
+    auto l2 = p.add_instruction(op, l1);
+    p.add_instruction(migraphx::op::squeeze{{1, 2}}, l2);
     auto prog = optimize_tf("mean_test_nhwc.pb", true);
 
     EXPECT(p == prog);
@@ -291,6 +308,23 @@ TEST_CASE(mul_test)
     EXPECT(p == prog);
 }
 
+TEST_CASE(onehot_test)
+{
+    migraphx::program p;
+    auto l0 = p.add_literal(
+        migraphx::literal{migraphx::shape{migraphx::shape::int32_type, {5}}, {1, 1, 1, 1, 1}});
+    p.add_literal(2);
+    p.add_literal(1.0f);
+    p.add_literal(0.0f);
+    auto l1 = p.add_literal(
+        migraphx::literal{migraphx::shape{migraphx::shape::float_type, {2, 2}}, {1, 0, 0, 1}});
+    int axis = 0;
+    p.add_instruction(migraphx::op::gather{axis}, l1, l0);
+    auto prog = optimize_tf("onehot_test.pb", false);
+
+    EXPECT(p == prog);
+}
+
 TEST_CASE(pack_test)
 {
     migraphx::program p;
@@ -475,20 +509,44 @@ TEST_CASE(stridedslice_test)
 {
     migraphx::program p;
     auto l0 = p.add_parameter("0", migraphx::shape{migraphx::shape::float_type, {1, 10, 1, 1}});
+    auto l1 = p.add_instruction(migraphx::op::transpose{{0, 2, 3, 1}}, l0);
     std::size_t num_axes = 4;
     migraphx::op::slice op;
     op.starts = {0, 0, 0, 0};
     op.ends   = {1, 1, 1, 5};
     op.axes   = std::vector<int64_t>(num_axes);
     std::iota(op.axes.begin(), op.axes.end(), 0);
-    auto l1          = p.add_instruction(op, l0);
+    auto l2          = p.add_instruction(op, l1);
     auto shrink_axis = 1;
-    p.add_instruction(migraphx::op::squeeze{{shrink_axis}}, l1);
+    p.add_instruction(migraphx::op::squeeze{{shrink_axis}}, l2);
     auto prog = optimize_tf("stridedslice_test.pb", true);
 
     EXPECT(p == prog);
 }
 
+TEST_CASE(stridedslice_masks_test)
+{
+    migraphx::program p;
+    auto l0 = p.add_parameter("0", migraphx::shape{migraphx::shape::float_type, {1, 10, 3, 3}});
+    std::size_t num_axes = 4;
+    migraphx::op::slice op;
+    op.starts = {0, 1, 1, 0};
+    op.ends   = {1, 3, 3, 10};
+    op.axes   = std::vector<int64_t>(num_axes);
+    std::iota(op.axes.begin(), op.axes.end(), 0);
+    // add literals for starts, ends, and strides in tf (NHWC format)
+    p.add_literal(migraphx::shape{migraphx::shape::int32_type, {4}}, std::vector<int>{0, 1, 1, 0});
+    p.add_literal(migraphx::shape{migraphx::shape::int32_type, {4}}, std::vector<int>{0, 0, 0, 0});
+    p.add_literal(migraphx::shape{migraphx::shape::int32_type, {4}}, std::vector<int>{1, 1, 1, 1});
+
+    auto l1 = p.add_instruction(migraphx::op::transpose{{0, 2, 3, 1}}, l0);
+    auto l2 = p.add_instruction(op, l1);
+    p.add_instruction(migraphx::op::transpose{{0, 3, 1, 2}}, l2);
+    auto prog = migraphx::parse_tf("stridedslice_masks_test.pb", true);
+
+    EXPECT(p == prog);
+}
+
 TEST_CASE(sub_test)
 {
     migraphx::program p;