Merge branch 'develop' of github.com:ROCmSoftwarePlatform/AMDMIGraphX into dyn_shape_update

test/api/test_array_base.cpp

+#include <migraphx/migraphx.hpp>
+#include "test.hpp"
+
+struct array2 : migraphx::array_base<array2>
+{
+    std::vector<int> v;
+    array2() = default;
+    array2(std::initializer_list<int> x) : v(x) {}
+    std::size_t size() const { return v.size(); }
+    int operator[](std::size_t i) const { return v[i]; }
+};
+
+TEST_CASE(iterators)
+{
+    array2 a = {1, 2, 3};
+    EXPECT(bool{std::equal(a.begin(), a.end(), a.v.begin())});
+}
+
+TEST_CASE(front_back)
+{
+    array2 a = {1, 2, 3};
+    EXPECT(a.front() == 1);
+    EXPECT(a.back() == 3);
+}
+
+TEST_CASE(empty)
+{
+    array2 a = {1, 2, 3};
+    EXPECT(not a.empty());
+}
+
+int main(int argc, const char* argv[]) { test::run(argc, argv); }

test/api/test_gpu.cpp

 #include <numeric>
+#include <hip/hip_runtime_api.h>
 #include <migraphx/migraphx.h>
 #include <migraphx/migraphx.hpp>
 #include "test.hpp"
@@ -38,6 +39,7 @@ TEST_CASE(load_and_run_ctx)
         pp.add(name, migraphx::argument::generate(param_shapes[name]));
     }
     auto ctx = p.experimental_get_context();
+    EXPECT(ctx.get_queue<hipStream_t>() != nullptr);
     p.eval(pp);
     ctx.finish();
 }

test/api/test_module_construct.cpp

@@ -60,16 +60,16 @@ TEST_CASE(if_then_else_op)
         p.compile(migraphx::target("ref"));
         auto outputs =
             p.eval({{"cond", migraphx::argument(cond_s, &cond)}, {"x", x_arg}, {"y", y_arg}});
-        return outputs;
+        return outputs[0];
     };
 
     // then branch
     auto then_res = run_prog(true);
-    CHECK(bool{then_res[0] == x_arg});
+    CHECK(bool{then_res == x_arg});
 
     // else branch
     auto else_res = run_prog(false);
-    CHECK(bool{else_res[0] == y_arg});
+    CHECK(bool{else_res == y_arg});
 }
 
 int main(int argc, const char* argv[]) { test::run(argc, argv); }

test/onnx/gathernd_batch_dims_test.onnx

+gathernd_batch_dims_test:

+/
+data
+indices
y"GatherND*
+
+batch_dims

gathernd_batch_dims_testZ
+data
+

+
+
+Z
+indices
+
+
+
b
+
y
+

+
+B
\ No newline at end of file

test/onnx/gathernd_test.onnx

+
gathernd_test:q
+
+data
+indices
y"GatherND
gathernd_testZ
+data
+

+
+Z
+indices
+
+
+b
+
y
+
+

+B
\ No newline at end of file

test/onnx/gen_onnx.py

@@ -1666,6 +1666,35 @@ def gather_elements_axis1_test():
     return ([node], [x, i], [y])
 
 
+@onnx_test
+def gathernd_test():
+    x = helper.make_tensor_value_info('data', TensorProto.FLOAT, [2, 2])
+    i = helper.make_tensor_value_info('indices', TensorProto.INT64, [2, 2])
+    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [2])
+
+    node = onnx.helper.make_node('GatherND',
+                                 inputs=['data', 'indices'],
+                                 outputs=['y'])
+
+    return ([node], [x, i], [y])
+
+
+@onnx_test
+def gathernd_batch_dims_test():
+    x = helper.make_tensor_value_info('data', TensorProto.FLOAT, [2, 2, 2])
+    i = helper.make_tensor_value_info('indices', TensorProto.INT64, [2, 1])
+    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [2, 2])
+
+    node = onnx.helper.make_node(
+        'GatherND',
+        inputs=['data', 'indices'],
+        outputs=['y'],
+        batch_dims=1,
+    )
+
+    return ([node], [x, i], [y])
+
+
 @onnx_test
 def gemm_test():
     x = helper.make_tensor_value_info('0', TensorProto.FLOAT, [5, 7])
@@ -4385,6 +4414,142 @@ def resize_upsample_pc_test():
     return ([node], [X], [Y], [scale_tensor])
 
 
+@onnx_test
+def reversesequence_4D_test():
+    x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [2, 2, 2, 2])
+    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [2, 2, 2, 2])
+
+    node = onnx.helper.make_node(
+        'ReverseSequence',
+        inputs=['x'],
+        outputs=['y'],
+        time_axis=0,
+        batch_axis=1,
+        sequence_lens=[2, 1],
+    )
+    return ([node], [x], [y])
+
+
+@onnx_test
+def reversesequence_batch_test():
+    x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [4, 4])
+    seq_lens = np.array([1, 2, 3, 4])
+    seq_lens_tensor = helper.make_tensor(
+        name="sequence_lens",
+        data_type=TensorProto.INT64,
+        dims=seq_lens.shape,
+        vals=seq_lens.astype(np.int64),
+    )
+    arg_seq_lens = helper.make_node(
+        "Constant",
+        inputs=[],
+        outputs=['arg_seq_lens'],
+        value=seq_lens_tensor,
+    )
+    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [4, 4])
+
+    node = onnx.helper.make_node(
+        'ReverseSequence',
+        inputs=['x', 'arg_seq_lens'],
+        outputs=['y'],
+        time_axis=1,
+        batch_axis=0,
+    )
+    return ([arg_seq_lens, node], [x], [y])
+
+
+@onnx_test
+def reversesequence_batch_axis_err_test():
+    x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [4, 4, 2])
+    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [4, 4, 2])
+
+    node = onnx.helper.make_node(
+        'ReverseSequence',
+        inputs=['x'],
+        outputs=['y'],
+        time_axis=0,
+        batch_axis=2,
+        sequence_lens=[4, 3, 2, 1],
+    )
+    return ([node], [x], [y])
+
+
+@onnx_test
+def reversesequence_rank_err_test():
+    x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [4])
+    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [4])
+
+    node = onnx.helper.make_node(
+        'ReverseSequence',
+        inputs=['x'],
+        outputs=['y'],
+        sequence_lens=[4, 3, 2, 1],
+    )
+    return ([node], [x], [y])
+
+
+@onnx_test
+def reversesequence_sequence_lens_shape_err_test():
+    x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [4, 4])
+    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [4, 4])
+
+    node = onnx.helper.make_node(
+        'ReverseSequence',
+        inputs=['x'],
+        outputs=['y'],
+        sequence_lens=[4, 3, 2],
+    )
+    return ([node], [x], [y])
+
+
+@onnx_test
+def reversesequence_same_axis_err_test():
+    x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [4, 4])
+    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [4, 4])
+
+    node = onnx.helper.make_node(
+        'ReverseSequence',
+        inputs=['x'],
+        outputs=['y'],
+        time_axis=1,
+        batch_axis=1,
+        sequence_lens=[4, 3, 2, 1],
+    )
+    return ([node], [x], [y])
+
+
+@onnx_test
+def reversesequence_time_axis_err_test():
+    x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [4, 4, 2, 3])
+    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [4, 4, 2, 3])
+
+    node = onnx.helper.make_node(
+        'ReverseSequence',
+        inputs=['x'],
+        outputs=['y'],
+        time_axis=3,
+        batch_axis=0,
+        sequence_lens=[4, 3, 2, 1],
+    )
+    return ([node], [x], [y])
+
+
+@onnx_test
+def reversesequence_time_test():
+    x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [4, 4])
+    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [4, 4])
+
+    node = onnx.helper.make_node(
+        'ReverseSequence',
+        inputs=['x'],
+        outputs=['y'],
+        time_axis=0,
+        batch_axis=1,
+        sequence_lens=[4, 3, 2, 1],
+    )
+    return ([node], [x], [y])
+
+
 @onnx_test
 def roialign_default_test():
     x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [10, 4, 7, 8])

test/onnx/onnx_test.cpp

@@ -1582,6 +1582,31 @@ TEST_CASE(gather_elements_axis1_test)
     EXPECT(p == prog);
 }
 
+TEST_CASE(gathernd_test)
+{
+    migraphx::program p;
+    auto* mm = p.get_main_module();
+    auto l0  = mm->add_parameter("data", migraphx::shape{migraphx::shape::float_type, {2, 2}});
+    auto l1  = mm->add_parameter("indices", migraphx::shape{migraphx::shape::int64_type, {2, 2}});
+    mm->add_instruction(migraphx::make_op("gathernd"), l0, l1);
+    auto prog = optimize_onnx("gathernd_test.onnx");
+
+    EXPECT(p == prog);
+}
+
+TEST_CASE(gathernd_batch_dims_test)
+{
+    migraphx::program p;
+    auto* mm = p.get_main_module();
+    auto l0  = mm->add_parameter("data", migraphx::shape{migraphx::shape::float_type, {2, 2, 2}});
+    auto l1  = mm->add_parameter("indices", migraphx::shape{migraphx::shape::int64_type, {2, 1}});
+    int batch_dims = 1;
+    mm->add_instruction(migraphx::make_op("gathernd", {{"batch_dims", batch_dims}}), l0, l1);
+    auto prog = optimize_onnx("gathernd_batch_dims_test.onnx");
+
+    EXPECT(p == prog);
+}
+
 TEST_CASE(gemm_test)
 {
     migraphx::program p;
@@ -4222,6 +4247,126 @@ TEST_CASE(resize_upsample_pf_test)
     EXPECT(p == prog);
 }
 
+TEST_CASE(reversesequence_batch_test)
+{
+    migraphx::program p;
+    auto* mm = p.get_main_module();
+
+    int batch_axis = 0;
+    int time_axis  = 1;
+
+    migraphx::shape sx{migraphx::shape::float_type, {4, 4}};
+    auto input = mm->add_parameter("x", sx);
+
+    std::vector<int64_t> sequence_lens = {1, 2, 3, 4};
+    mm->add_literal({{migraphx::shape::int64_type, {4}}, sequence_lens});
+
+    int batch_size = sx.lens()[batch_axis];
+    int time_size  = sx.lens()[time_axis];
+
+    auto add_slice =
+        [&mm, &input, batch_axis, time_axis](int b_start, int b_end, int t_start, int t_end) {
+            return mm->add_instruction(migraphx::make_op("slice",
+                                                         {{"axes", {batch_axis, time_axis}},
+                                                          {"starts", {b_start, t_start}},
+                                                          {"ends", {b_end, t_end}}}),
+                                       input);
+        };
+    auto ret = add_slice(0, 1, 0, time_size);
+    for(int b = 1; b < batch_size; ++b)
+    {
+        auto s0 = add_slice(b, b + 1, 0, sequence_lens[b]);
+        s0      = mm->add_instruction(migraphx::make_op("reverse", {{"axes", {time_axis}}}), s0);
+        if(sequence_lens[b] < time_size)
+        {
+            auto s1 = add_slice(b, b + 1, sequence_lens[b], time_size);
+            s0 = mm->add_instruction(migraphx::make_op("concat", {{"axis", time_axis}}), s0, s1);
+        }
+        ret = mm->add_instruction(migraphx::make_op("concat", {{"axis", batch_axis}}), ret, s0);
+    }
+    mm->add_return({ret});
+
+    auto prog = migraphx::parse_onnx("reversesequence_batch_test.onnx");
+    EXPECT(p == prog);
+}
+
+TEST_CASE(reversesequence_batch_axis_err_test)
+{
+    EXPECT(test::throws([&] { migraphx::parse_onnx("reversesequence_batch_axis_err_test.onnx"); }));
+}
+
+TEST_CASE(reversesequence_rank_err_test)
+{
+    EXPECT(test::throws([&] { migraphx::parse_onnx("reversesequence_rank_err_test.onnx"); }));
+}
+
+TEST_CASE(reversesequence_sequence_lens_shape_err_test)
+{
+    EXPECT(test::throws(
+        [&] { migraphx::parse_onnx("reversesequence_sequence_lens_shape_err_test.onnx"); }));
+}
+
+TEST_CASE(reversesequence_same_axis_err_test)
+{
+    EXPECT(test::throws([&] { migraphx::parse_onnx("reversesequence_same_axis_err_test.onnx"); }));
+}
+
+TEST_CASE(reversesequence_time_axis_err_test)
+{
+    EXPECT(test::throws([&] { migraphx::parse_onnx("reversesequence_time_axis_err_test.onnx"); }));
+}
+
+TEST_CASE(reversesequence_time_test)
+{
+    migraphx::program p;
+    auto* mm = p.get_main_module();
+
+    int batch_axis = 1;
+    int time_axis  = 0;
+
+    migraphx::shape sx{migraphx::shape::float_type, {4, 4}};
+    auto input = mm->add_parameter("x", sx);
+
+    int batch_size                     = sx.lens()[batch_axis];
+    int time_size                      = sx.lens()[time_axis];
+    std::vector<int64_t> sequence_lens = {4, 3, 2, 1};
+
+    auto add_slice =
+        [&mm, &input, batch_axis, time_axis](int b_start, int b_end, int t_start, int t_end) {
+            return mm->add_instruction(migraphx::make_op("slice",
+                                                         {{"axes", {batch_axis, time_axis}},
+                                                          {"starts", {b_start, t_start}},
+                                                          {"ends", {b_end, t_end}}}),
+                                       input);
+        };
+
+    migraphx::instruction_ref ret;
+    for(int b = 0; b < batch_size - 1; ++b)
+    {
+        auto s0 = add_slice(b, b + 1, 0, sequence_lens[b]);
+        s0      = mm->add_instruction(migraphx::make_op("reverse", {{"axes", {time_axis}}}), s0);
+        if(sequence_lens[b] < time_size)
+        {
+            auto s1 = add_slice(b, b + 1, sequence_lens[b], time_size);
+            s0 = mm->add_instruction(migraphx::make_op("concat", {{"axis", time_axis}}), s0, s1);
+        }
+        if(b == 0)
+        {
+            ret = s0;
+        }
+        else
+        {
+            ret = mm->add_instruction(migraphx::make_op("concat", {{"axis", batch_axis}}), ret, s0);
+        }
+    }
+    auto s0 = add_slice(batch_size - 1, batch_size, 0, time_size);
+    ret     = mm->add_instruction(migraphx::make_op("concat", {{"axis", batch_axis}}), ret, s0);
+    mm->add_return({ret});
+
+    auto prog = migraphx::parse_onnx("reversesequence_time_test.onnx");
+    EXPECT(p == prog);
+}
+
 TEST_CASE(roialign_default_test)
 {
     migraphx::shape sx{migraphx::shape::float_type, {10, 4, 7, 8}};

test/onnx/reversesequence_rank_err_test.onnx

+reversesequence_rank_err_test:v
+3
+
x
y"ReverseSequence*
+sequence_lens@@@@

reversesequence_rank_err_testZ
+
x
+
+

+b
+
y
+
+

+B
\ No newline at end of file

test/onnx/reversesequence_same_axis_err_test.onnx

+"reversesequence_same_axis_err_test:

+X
+
x
y"ReverseSequence*
+
+batch_axis

*
+sequence_lens@@@@

*
+	time_axis

"reversesequence_same_axis_err_testZ
+
x
+

+
+b
+
y
+

+
+B
\ No newline at end of file

test/onnx/reversesequence_sequence_lens_shape_err_test.onnx

+,reversesequence_sequence_lens_shape_err_test:‹

+1
+
x
y"ReverseSequence*
+sequence_lens@@@ 
,reversesequence_sequence_lens_shape_err_testZ
+
x
+

+
+b
+
y
+

+
+B
\ No newline at end of file

test/onnx/verify_onnx.cpp

@@ -698,6 +698,69 @@ TEST_CASE(resize_upsample_pf_test)
     EXPECT(migraphx::verify_range(result_vector, gold));
 }
 
+TEST_CASE(reversesequence_4D_verify_test)
+{
+    migraphx::program p = migraphx::parse_onnx("reversesequence_4D_test.onnx");
+    p.compile(migraphx::ref::target{});
+
+    migraphx::shape xs{migraphx::shape::float_type, {2, 2, 2, 2}};
+    std::vector<float> x_data = {
+        0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0};
+    migraphx::parameter_map param_map;
+    param_map["x"] = migraphx::argument(xs, x_data.data());
+
+    auto result = p.eval(param_map).back();
+    std::vector<float> result_vector;
+    result.visit([&](auto output) { result_vector.assign(output.begin(), output.end()); });
+
+    std::vector<float> gold = {
+        8.0, 9.0, 10.0, 11.0, 4.0, 5.0, 6.0, 7.0, 0.0, 1.0, 2.0, 3.0, 12.0, 13.0, 14.0, 15.0};
+
+    EXPECT(migraphx::verify_range(result_vector, gold));
+}
+
+TEST_CASE(reversesequence_batch_verify_test)
+{
+    migraphx::program p = migraphx::parse_onnx("reversesequence_batch_test.onnx");
+    p.compile(migraphx::ref::target{});
+
+    migraphx::shape xs{migraphx::shape::float_type, {4, 4}};
+    std::vector<float> x_data = {
+        0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0};
+    migraphx::parameter_map param_map;
+    param_map["x"] = migraphx::argument(xs, x_data.data());
+
+    auto result = p.eval(param_map).back();
+    std::vector<float> result_vector;
+    result.visit([&](auto output) { result_vector.assign(output.begin(), output.end()); });
+
+    std::vector<float> gold = {
+        0.0, 1.0, 2.0, 3.0, 5.0, 4.0, 6.0, 7.0, 10.0, 9.0, 8.0, 11.0, 15.0, 14.0, 13.0, 12.0};
+
+    EXPECT(migraphx::verify_range(result_vector, gold));
+}
+
+TEST_CASE(reversesequence_time_verify_test)
+{
+    migraphx::program p = migraphx::parse_onnx("reversesequence_time_test.onnx");
+    p.compile(migraphx::ref::target{});
+
+    migraphx::shape xs{migraphx::shape::float_type, {4, 4}};
+    std::vector<float> x_data = {
+        0.0, 4.0, 8.0, 12.0, 1.0, 5.0, 9.0, 13.0, 2.0, 6.0, 10.0, 14.0, 3.0, 7.0, 11.0, 15.0};
+    migraphx::parameter_map param_map;
+    param_map["x"] = migraphx::argument(xs, x_data.data());
+
+    auto result = p.eval(param_map).back();
+    std::vector<float> result_vector;
+    result.visit([&](auto output) { result_vector.assign(output.begin(), output.end()); });
+
+    std::vector<float> gold = {
+        3.0, 6.0, 9.0, 12.0, 2.0, 5.0, 8.0, 13.0, 1.0, 4.0, 10.0, 14.0, 0.0, 7.0, 11.0, 15.0};
+
+    EXPECT(migraphx::verify_range(result_vector, gold));
+}
+
 TEST_CASE(selu_test)
 {
     migraphx::program p = migraphx::parse_onnx("selu_test.onnx");

test/py/onnx_backend_test.py

@@ -178,6 +178,7 @@ def create_backend_test(testname=None, target_device=None):
         backend_test.include(r'.*test_reduce.*')
         backend_test.include(r'.*test_ReLU*')
         backend_test.include(r'.*test_relu.*')
+        #backend_test.include(r'.*test_reversesequence.*')
         backend_test.include(r'.*test_RoiAlign*')
         backend_test.include(r'.*test_roialign.*')
         backend_test.include(r'.*test_scatter.*')
@@ -267,9 +268,6 @@ def create_backend_test(testname=None, target_device=None):
         backend_test.exclude(r'test_expand_shape_model2_cpu')
         backend_test.exclude(r'test_expand_shape_model3_cpu')
         backend_test.exclude(r'test_expand_shape_model4_cpu')
-        backend_test.exclude(r'test_gathernd_example_float32_cpu')
-        backend_test.exclude(r'test_gathernd_example_int32_batch_dim1_cpu')
-        backend_test.exclude(r'test_gathernd_example_int32_cpu')
         backend_test.exclude(r'test_identity_sequence_cpu')
         backend_test.exclude(r'test_maxpool_2d_uint8_cpu')
         backend_test.exclude(r'test_negative_log_likelihood_loss_*')

test/ref_ops_test.cpp

@@ -1653,6 +1653,203 @@ TEST_CASE(gather_test)
     }
 }
 
+TEST_CASE(gathernd_test)
+{
+    {
+        migraphx::program p;
+        auto* mm = p.get_main_module();
+
+        migraphx::shape ds{migraphx::shape::float_type, {2, 2}};
+        migraphx::shape is{migraphx::shape::int64_type, {2, 2}};
+
+        std::vector<float> data_vec(2 * 2);
+        std::iota(data_vec.begin(), data_vec.end(), 0);
+        std::vector<int64_t> indices_vec{0, 0, 1, 1};
+
+        auto data    = mm->add_literal(migraphx::literal{ds, data_vec});
+        auto indices = mm->add_literal(migraphx::literal{is, indices_vec});
+
+        mm->add_instruction(migraphx::make_op("gathernd"), data, indices);
+        p.compile(migraphx::ref::target{});
+        auto result = p.eval({}).back();
+        std::vector<float> res_data{};
+        std::vector<float> gold{0, 3};
+        result.visit([&](auto output) { res_data.assign(output.begin(), output.end()); });
+
+        EXPECT(migraphx::verify_range(res_data, gold));
+    }
+
+    {
+        migraphx::program p;
+        auto* mm = p.get_main_module();
+
+        migraphx::shape ds{migraphx::shape::float_type, {2, 2}};
+        migraphx::shape is{migraphx::shape::int64_type, {2, 1}};
+
+        std::vector<float> data_vec(2 * 2);
+        std::iota(data_vec.begin(), data_vec.end(), 0);
+        std::vector<int64_t> indices_vec{1, 0};
+
+        auto data    = mm->add_literal(migraphx::literal{ds, data_vec});
+        auto indices = mm->add_literal(migraphx::literal{is, indices_vec});
+
+        mm->add_instruction(migraphx::make_op("gathernd"), data, indices);
+        p.compile(migraphx::ref::target{});
+        auto result = p.eval({}).back();
+        std::vector<float> res_data{};
+        std::vector<float> gold{2, 3, 0, 1};
+        result.visit([&](auto output) { res_data.assign(output.begin(), output.end()); });
+
+        EXPECT(migraphx::verify_range(res_data, gold));
+    }
+
+    {
+        migraphx::program p;
+        auto* mm = p.get_main_module();
+
+        migraphx::shape ds{migraphx::shape::float_type, {2, 3, 1}};
+        migraphx::shape is{migraphx::shape::int64_type, {2, 2, 1}};
+
+        std::vector<float> data_vec(2 * 3 * 1);
+        std::iota(data_vec.begin(), data_vec.end(), 0);
+        std::vector<int64_t> indices_vec{1, 0, 0, 1};
+
+        auto data    = mm->add_literal(migraphx::literal{ds, data_vec});
+        auto indices = mm->add_literal(migraphx::literal{is, indices_vec});
+
+        mm->add_instruction(migraphx::make_op("gathernd"), data, indices);
+        p.compile(migraphx::ref::target{});
+        auto result = p.eval({}).back();
+        std::vector<float> res_data{};
+        std::vector<float> gold{3, 4, 5, 0, 1, 2, 0, 1, 2, 3, 4, 5};
+        result.visit([&](auto output) { res_data.assign(output.begin(), output.end()); });
+
+        EXPECT(migraphx::verify_range(res_data, gold));
+    }
+
+    {
+        migraphx::program p;
+        auto* mm = p.get_main_module();
+
+        migraphx::shape ds{migraphx::shape::float_type, {2, 3, 2, 3}};
+        migraphx::shape is{migraphx::shape::int64_type, {2, 2, 2}};
+
+        std::vector<float> data_vec(2 * 3 * 2 * 3);
+        std::iota(data_vec.begin(), data_vec.end(), 0);
+        std::vector<int64_t> indices_vec{0, 0, 0, 1, 0, 0, 0, 1};
+        const int batch_dims = 1;
+
+        auto data    = mm->add_literal(migraphx::literal{ds, data_vec});
+        auto indices = mm->add_literal(migraphx::literal{is, indices_vec});
+
+        mm->add_instruction(
+            migraphx::make_op("gathernd", {{"batch_dims", batch_dims}}), data, indices);
+        p.compile(migraphx::ref::target{});
+        auto result = p.eval({}).back();
+        std::vector<float> res_data{};
+        std::vector<float> gold{0, 1, 2, 3, 4, 5, 18, 19, 20, 21, 22, 23};
+        result.visit([&](auto output) { res_data.assign(output.begin(), output.end()); });
+
+        EXPECT(migraphx::verify_range(res_data, gold));
+    }
+
+    {
+        migraphx::program p;
+        auto* mm = p.get_main_module();
+
+        migraphx::shape ds{migraphx::shape::float_type, {2, 3, 1, 3}};
+        migraphx::shape is{migraphx::shape::int64_type, {2, 3, 2}};
+
+        std::vector<float> data_vec(2 * 3 * 1 * 3);
+        std::iota(data_vec.begin(), data_vec.end(), 0);
+        std::vector<int64_t> indices_vec{0, 0, 0, 1, 0, 2, 0, 2, 0, 1, 0, 0};
+        const int batch_dims = 2;
+
+        auto data    = mm->add_literal(migraphx::literal{ds, data_vec});
+        auto indices = mm->add_literal(migraphx::literal{is, indices_vec});
+
+        mm->add_instruction(
+            migraphx::make_op("gathernd", {{"batch_dims", batch_dims}}), data, indices);
+
+        p.compile(migraphx::ref::target{});
+        auto result = p.eval({}).back();
+        std::vector<float> res_data{};
+        std::vector<float> gold{0, 4, 8, 11, 13, 15};
+        result.visit([&](auto output) { res_data.assign(output.begin(), output.end()); });
+
+        EXPECT(migraphx::verify_range(res_data, gold));
+    }
+
+    {
+        // k > r - batch_dims
+        migraphx::program p;
+        auto* mm = p.get_main_module();
+
+        migraphx::shape ds{migraphx::shape::float_type, {2, 3, 1, 3}};
+        migraphx::shape is{migraphx::shape::int64_type, {2, 3, 3}};
+
+        std::vector<float> data_vec(2 * 3 * 1 * 3);
+        std::iota(data_vec.begin(), data_vec.end(), 0);
+        std::vector<int64_t> indices_vec(2 * 3 * 3, 0);
+        const int batch_dims = 2;
+
+        auto data    = mm->add_literal(migraphx::literal{ds, data_vec});
+        auto indices = mm->add_literal(migraphx::literal{is, indices_vec});
+
+        EXPECT(test::throws([&] {
+            mm->add_instruction(
+                migraphx::make_op("gathernd", {{"batch_dims", batch_dims}}), data, indices);
+        }));
+    }
+}
+
+TEST_CASE(gathernd_negative_index_test)
+{
+    {
+        migraphx::program p;
+        auto* mm = p.get_main_module();
+
+        migraphx::shape ds{migraphx::shape::float_type, {2, 2}};
+        migraphx::shape is{migraphx::shape::int64_type, {2, 1, 1}};
+
+        std::vector<float> data_vec(2 * 2);
+        std::iota(data_vec.begin(), data_vec.end(), 0);
+        std::vector<int64_t> indices_vec{-1, 0};
+
+        auto data    = mm->add_literal(migraphx::literal{ds, data_vec});
+        auto indices = mm->add_literal(migraphx::literal{is, indices_vec});
+
+        mm->add_instruction(migraphx::make_op("gathernd"), data, indices);
+        p.compile(migraphx::ref::target{});
+        auto result = p.eval({}).back();
+        std::vector<float> res_data{};
+        std::vector<float> gold{2, 3, 0, 1};
+        result.visit([&](auto output) { res_data.assign(output.begin(), output.end()); });
+
+        EXPECT(migraphx::verify_range(res_data, gold));
+    }
+
+    {
+        migraphx::program p;
+        auto* mm = p.get_main_module();
+
+        migraphx::shape ds{migraphx::shape::float_type, {2, 2}};
+        migraphx::shape is{migraphx::shape::int64_type, {2, 1, 1}};
+
+        std::vector<float> data_vec(2 * 2);
+        std::iota(data_vec.begin(), data_vec.end(), 0);
+        std::vector<int64_t> indices_vec{-3, 0};
+
+        auto data    = mm->add_literal(migraphx::literal{ds, data_vec});
+        auto indices = mm->add_literal(migraphx::literal{is, indices_vec});
+
+        mm->add_instruction(migraphx::make_op("gathernd"), data, indices);
+        p.compile(migraphx::ref::target{});
+
+        EXPECT(test::throws([&] { p.eval({}); }));
+    }
+}
+
 TEST_CASE(globalavgpool_test)
 {
     migraphx::program p;

test/verify/test_gathernd_batch_dims_1.cpp

+
+#include "verify_program.hpp"
+#include <migraphx/program.hpp>
+#include <migraphx/generate.hpp>
+#include <migraphx/make_op.hpp>
+
+struct test_gathernd_batch_dims_1 : verify_program<test_gathernd_batch_dims_1>
+{
+    migraphx::program create_program() const
+    {
+        migraphx::program p;
+        auto* mm = p.get_main_module();
+        migraphx::shape ds{migraphx::shape::float_type, {2, 3, 2, 3}};
+        migraphx::shape is{migraphx::shape::int64_type, {2, 3, 2}};
+        std::vector<int64_t> indices{1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0};
+        auto a0        = mm->add_parameter("data", ds);
+        auto a1        = mm->add_literal(migraphx::literal{is, indices});
+        int batch_dims = 1;
+        mm->add_instruction(migraphx::make_op("gathernd", {{"batch_dims", batch_dims}}), a0, a1);
+        return p;
+    }
+};

test/verify/test_gathernd_batch_dims_2.cpp

+#include "verify_program.hpp"
+#include <migraphx/program.hpp>
+#include <migraphx/generate.hpp>
+#include <migraphx/make_op.hpp>
+
+struct test_gathernd_batch_dims_2 : verify_program<test_gathernd_batch_dims_2>
+{
+    migraphx::program create_program() const
+    {
+        migraphx::program p;
+        auto* mm = p.get_main_module();
+        migraphx::shape ds{migraphx::shape::float_type, {2, 3, 1, 3}};
+        migraphx::shape is{migraphx::shape::int64_type, {2, 3, 2}};
+        std::vector<int64_t> indices{0, 0, 0, 1, 0, 2, 0, 2, 0, 1, 0, 0};
+        auto a0        = mm->add_parameter("data", ds);
+        auto a1        = mm->add_literal(migraphx::literal{is, indices});
+        int batch_dims = 2;
+        mm->add_instruction(migraphx::make_op("gathernd", {{"batch_dims", batch_dims}}), a0, a1);
+        return p;
+    }
+};