diff --git a/src/plugins/intel_gpu/include/intel_gpu/primitives/gather_elements.hpp b/src/plugins/intel_gpu/include/intel_gpu/primitives/gather_elements.hpp index 9134f396a78..d23439f1e7b 100644 --- a/src/plugins/intel_gpu/include/intel_gpu/primitives/gather_elements.hpp +++ b/src/plugins/intel_gpu/include/intel_gpu/primitives/gather_elements.hpp @@ -19,15 +19,6 @@ namespace cldnn { struct gather_elements : public primitive_base { CLDNN_DECLARE_PRIMITIVE(gather_elements) - enum gather_elements_axis { - along_b, - along_f, - along_x, - along_y, - along_z, - along_w - }; - /// @brief Constructs gather_elements primitive. /// @param id This primitive id. /// @param data Input data primitive id. @@ -40,7 +31,7 @@ struct gather_elements : public primitive_base { const primitive_id& indices, const format& output_format, const tensor& output_shape, - const gather_elements_axis axis, + const int64_t axis, const primitive_id& ext_prim_id = "", const padding& output_padding = padding()) : primitive_base(id, {data, indices}, ext_prim_id, output_padding), output_format(output_format), output_shape(output_shape), axis(axis) {} @@ -51,7 +42,7 @@ struct gather_elements : public primitive_base { tensor output_shape; /// @brief Which axis to gather on. - gather_elements_axis axis; + int64_t axis; }; /// @} /// @} diff --git a/src/plugins/intel_gpu/src/graph/impls/ocl/gather_elements.cpp b/src/plugins/intel_gpu/src/graph/impls/ocl/gather_elements.cpp index 7c8a10bc22e..b1c9d72c155 100644 --- a/src/plugins/intel_gpu/src/graph/impls/ocl/gather_elements.cpp +++ b/src/plugins/intel_gpu/src/graph/impls/ocl/gather_elements.cpp @@ -14,22 +14,35 @@ using namespace cldnn; namespace cldnn { namespace ocl { -kernel_selector::gather_elements_axis convert_axis(gather_elements::gather_elements_axis axis) { + +static inline kernel_selector::gather_elements_axis convert_axis(int64_t axis, size_t rank) { + if (axis < 0) { + axis += rank; + } switch (axis) { - case gather_elements::along_x: - return kernel_selector::gather_elements_axis::X; - case gather_elements::along_y: - return kernel_selector::gather_elements_axis::Y; - case gather_elements::along_z: - return kernel_selector::gather_elements_axis::Z; - case gather_elements::along_w: - return kernel_selector::gather_elements_axis::W; - case gather_elements::along_f: - return kernel_selector::gather_elements_axis::FEATURE; - case gather_elements::along_b: - return kernel_selector::gather_elements_axis::BATCH; - default: - return kernel_selector::gather_elements_axis::BATCH; + case 0: return kernel_selector::gather_elements_axis::BATCH; + case 1: return kernel_selector::gather_elements_axis::FEATURE; + case 2: + if (rank == 6) + return kernel_selector::gather_elements_axis::W; + else if (rank == 5) + return kernel_selector::gather_elements_axis::Z; + else + return kernel_selector::gather_elements_axis::Y; + case 3: + if (rank == 6) + return kernel_selector::gather_elements_axis::Z; + else if (rank == 5) + return kernel_selector::gather_elements_axis::Y; + else + return kernel_selector::gather_elements_axis::X; + case 4: + if (rank == 6) + return kernel_selector::gather_elements_axis::Y; + else + return kernel_selector::gather_elements_axis::X; + case 5: return kernel_selector::gather_elements_axis::X; + default: IE_THROW() << "Incorrect gather_elements axis."; } } @@ -48,7 +61,8 @@ public: auto gather_elements_optional_params = get_default_optional_params(arg.get_program()); - gather_elements_params.axis = convert_axis(prim->axis); + size_t rank = arg.get_output_layout().get_rank(); + gather_elements_params.axis = convert_axis(prim->axis, rank); gather_elements_params.inputs.push_back(convert_data_tensor(impl_param.input_layouts[1])); diff --git a/src/plugins/intel_gpu/src/plugin/ops/gather_elements.cpp b/src/plugins/intel_gpu/src/plugin/ops/gather_elements.cpp index fa5c0c5e93f..68ec21c2289 100644 --- a/src/plugins/intel_gpu/src/plugin/ops/gather_elements.cpp +++ b/src/plugins/intel_gpu/src/plugin/ops/gather_elements.cpp @@ -13,50 +13,26 @@ namespace ov { namespace intel_gpu { -static cldnn::gather_elements::gather_elements_axis GetGatherAxis(int axis, unsigned rank) { - if (axis < 0) - axis += rank; - if (axis < 0 || axis >= static_cast(rank)) - IE_THROW() << "GatherElements axis is not correspond to number of dimensions"; - - // Difference in dimension ordering between IE and GPU plugin, - // reverse spatial dimensions after batch and feature. - unsigned cldnn_axis = axis; - if (axis >= 2) { - auto spatial_axis = axis - 2; - // Default and minimum number of dimensions is 4 - auto spatial_size = std::max(rank, 4u) - 2; - cldnn_axis = spatial_size - spatial_axis - 1 + 2; - } - - switch (cldnn_axis) { - case 0: return cldnn::gather_elements::gather_elements_axis::along_b; - case 1: return cldnn::gather_elements::gather_elements_axis::along_f; - case 2: return cldnn::gather_elements::gather_elements_axis::along_x; - case 3: return cldnn::gather_elements::gather_elements_axis::along_y; - case 4: return cldnn::gather_elements::gather_elements_axis::along_z; - case 5: return cldnn::gather_elements::gather_elements_axis::along_w; - default: IE_THROW() << "Unsupported GatherElements axis: " << axis; - } - return cldnn::gather_elements::gather_elements_axis::along_f; // shouldn't get here -} - static void CreateGatherElementsOp(Program& p, const std::shared_ptr& op) { p.ValidateInputs(op, {2}); auto inputPrimitives = p.GetInputPrimitiveIDs(op); std::string layerName = layer_type_name_ID(op); - size_t rank = op->get_input_shape(0).size(); - int32_t axis = static_cast(op->get_axis()); - auto outLayout = cldnn::format::get_default_format(op->get_output_shape(0).size()); + size_t rank = op->get_input_shape(0).size(); + int64_t axis = op->get_axis(); + if (axis < 0) + axis += rank; + OPENVINO_ASSERT(axis >= 0 && axis < static_cast(rank), + "GatherElements axis is not correspond to number of dimensions"); + auto primitive = cldnn::gather_elements(layerName, inputPrimitives[0], inputPrimitives[1], outLayout, tensor_from_dims(op->get_output_shape(0)), - GetGatherAxis(axis, rank), + axis, op->get_friendly_name()); p.AddPrimitive(primitive); diff --git a/src/plugins/intel_gpu/tests/fusions/gather_elements_fusion_test.cpp b/src/plugins/intel_gpu/tests/fusions/gather_elements_fusion_test.cpp index 7b4c1bedc41..a921365ff46 100644 --- a/src/plugins/intel_gpu/tests/fusions/gather_elements_fusion_test.cpp +++ b/src/plugins/intel_gpu/tests/fusions/gather_elements_fusion_test.cpp @@ -29,7 +29,7 @@ struct gather_elements_test_params { format output_format; tensor output_shape; - cldnn::gather_elements::gather_elements_axis axis; + int64_t axis; data_types default_type; format default_format; @@ -50,22 +50,10 @@ public: } size_t get_axis_dim(gather_elements_test_params& p) { - switch (p.axis) { - case cldnn::gather_elements::gather_elements_axis::along_x: - return p.input_shape.spatial[0]; - case cldnn::gather_elements::gather_elements_axis::along_y: - return p.input_shape.spatial[1]; - case cldnn::gather_elements::gather_elements_axis::along_z: - return p.input_shape.spatial[2]; - case cldnn::gather_elements::gather_elements_axis::along_w: - return p.input_shape.spatial[3]; - case cldnn::gather_elements::gather_elements_axis::along_f: - return p.input_shape.feature[0]; - case cldnn::gather_elements::gather_elements_axis::along_b: - return p.input_shape.batch[0]; - default: - return 1; - } + size_t rank = p.input_format.dimension(); + auto default_fmt = format::get_default_format(rank); + auto dims = p.input_shape.sizes(default_fmt); + return dims[p.axis]; } layout get_input_layout(gather_elements_test_params& p) { @@ -89,28 +77,28 @@ public: /* ----------------------------------------------------------------------------------------------------- */ /* ------------------------------------------ GatherElements cases ------------------------------------- */ /* ----------------------------------------------------------------------------------------------------- */ -#define CASE_GATHER_ELEMENTS_FP16_4D_1 data_types::f16, format::bfyx, { 3, 7, 9, 8 }, format::bfyx, { 3, 7, 9, 8 }, format::bfyx, { 3, 7, 9, 8 }, cldnn::gather_elements::gather_elements_axis::along_y, data_types::f16, format::bfyx -#define CASE_GATHER_ELEMENTS_FP16_4D_2 data_types::f16, format::bfyx, { 3, 2, 8, 3 }, format::bfyx, { 2, 2, 8, 3 }, format::bfyx, { 2, 2, 8, 3 }, cldnn::gather_elements::gather_elements_axis::along_b, data_types::f16, format::bfyx -#define CASE_GATHER_ELEMENTS_FP16_4D_3 data_types::f16, format::bfyx, { 1, 3, 2, 9 }, format::bfyx, { 1, 3, 5, 9 }, format::bfyx, { 1, 3, 5, 9 }, cldnn::gather_elements::gather_elements_axis::along_x, data_types::f16, format::bfyx +#define CASE_GATHER_ELEMENTS_FP16_4D_1 data_types::f16, format::bfyx, { 3, 7, 9, 8 }, format::bfyx, { 3, 7, 9, 8 }, format::bfyx, { 3, 7, 9, 8 }, 2, data_types::f16, format::bfyx +#define CASE_GATHER_ELEMENTS_FP16_4D_2 data_types::f16, format::bfyx, { 3, 2, 8, 3 }, format::bfyx, { 2, 2, 8, 3 }, format::bfyx, { 2, 2, 8, 3 }, 0, data_types::f16, format::bfyx +#define CASE_GATHER_ELEMENTS_FP16_4D_3 data_types::f16, format::bfyx, { 1, 3, 2, 9 }, format::bfyx, { 1, 3, 5, 9 }, format::bfyx, { 1, 3, 5, 9 }, 3, data_types::f16, format::bfyx -#define CASE_GATHER_ELEMENTS_FP16_5D_1 data_types::f16, format::bfzyx, { 3, 2, 5, 2, 3 }, format::bfzyx, { 3, 2, 2, 2, 3 }, format::bfzyx, { 3, 2, 2, 2, 3 }, cldnn::gather_elements::gather_elements_axis::along_x, data_types::f16, format::bfzyx -#define CASE_GATHER_ELEMENTS_FP16_5D_2 data_types::f16, format::bfzyx, { 5, 4, 7, 4, 4 }, format::bfzyx, { 5, 4, 7, 4, 3 }, format::bfzyx, { 5, 4, 7, 4, 3 }, cldnn::gather_elements::gather_elements_axis::along_z, data_types::f16, format::bfzyx +#define CASE_GATHER_ELEMENTS_FP16_5D_1 data_types::f16, format::bfzyx, { 3, 2, 5, 2, 3 }, format::bfzyx, { 3, 2, 2, 2, 3 }, format::bfzyx, { 3, 2, 2, 2, 3 }, 4, data_types::f16, format::bfzyx +#define CASE_GATHER_ELEMENTS_FP16_5D_2 data_types::f16, format::bfzyx, { 5, 4, 7, 4, 4 }, format::bfzyx, { 5, 4, 7, 4, 3 }, format::bfzyx, { 5, 4, 7, 4, 3 }, 2, data_types::f16, format::bfzyx -#define CASE_GATHER_ELEMENTS_FP16_6D_1 data_types::f16, format::bfwzyx, { 5, 4, 6, 7, 8, 2 }, format::bfwzyx, { 5, 2, 6, 7, 8, 2 }, format::bfwzyx, { 5, 2, 6, 7, 8, 2 }, cldnn::gather_elements::gather_elements_axis::along_f, data_types::f16, format::bfwzyx -#define CASE_GATHER_ELEMENTS_FP16_6D_2 data_types::f16, format::bfwzyx, { 2, 1, 2, 3, 2, 1 }, format::bfwzyx, { 2, 1, 2, 3, 2, 3 }, format::bfwzyx, { 2, 1, 2, 3, 2, 3 }, cldnn::gather_elements::gather_elements_axis::along_w, data_types::f16, format::bfwzyx -#define CASE_GATHER_ELEMENTS_FP16_6D_3 data_types::f16, format::bfwzyx, { 2, 2, 3, 4, 4, 2 }, format::bfwzyx, { 2, 2, 6, 4, 4, 2 }, format::bfwzyx, { 2, 2, 6, 4, 4, 2 }, cldnn::gather_elements::gather_elements_axis::along_x, data_types::f16, format::bfwzyx +#define CASE_GATHER_ELEMENTS_FP16_6D_1 data_types::f16, format::bfwzyx, { 5, 4, 6, 7, 8, 2 }, format::bfwzyx, { 5, 2, 6, 7, 8, 2 }, format::bfwzyx, { 5, 2, 6, 7, 8, 2 }, 1, data_types::f16, format::bfwzyx +#define CASE_GATHER_ELEMENTS_FP16_6D_2 data_types::f16, format::bfwzyx, { 2, 1, 2, 3, 2, 1 }, format::bfwzyx, { 2, 1, 2, 3, 2, 3 }, format::bfwzyx, { 2, 1, 2, 3, 2, 3 }, 2, data_types::f16, format::bfwzyx +#define CASE_GATHER_ELEMENTS_FP16_6D_3 data_types::f16, format::bfwzyx, { 2, 2, 3, 4, 4, 2 }, format::bfwzyx, { 2, 2, 6, 4, 4, 2 }, format::bfwzyx, { 2, 2, 6, 4, 4, 2 }, 5, data_types::f16, format::bfwzyx -#define CASE_GATHER_ELEMENTS_FP32_4D_1 data_types::f32, format::bfyx, { 3, 7, 9, 8 }, format::bfyx, { 3, 7, 9, 8 }, format::bfyx, { 3, 7, 9, 8 }, cldnn::gather_elements::gather_elements_axis::along_y, data_types::f32, format::bfyx -#define CASE_GATHER_ELEMENTS_FP32_4D_2 data_types::f32, format::bfyx, { 3, 2, 8, 3 }, format::bfyx, { 2, 2, 8, 3 }, format::bfyx, { 2, 2, 8, 3 }, cldnn::gather_elements::gather_elements_axis::along_b, data_types::f32, format::bfyx -#define CASE_GATHER_ELEMENTS_FP32_4D_3 data_types::f32, format::bfyx, { 1, 3, 2, 9 }, format::bfyx, { 1, 3, 5, 9 }, format::bfyx, { 1, 3, 5, 9 }, cldnn::gather_elements::gather_elements_axis::along_x, data_types::f32, format::bfyx +#define CASE_GATHER_ELEMENTS_FP32_4D_1 data_types::f32, format::bfyx, { 3, 7, 9, 8 }, format::bfyx, { 3, 7, 9, 8 }, format::bfyx, { 3, 7, 9, 8 }, 2, data_types::f32, format::bfyx +#define CASE_GATHER_ELEMENTS_FP32_4D_2 data_types::f32, format::bfyx, { 3, 2, 8, 3 }, format::bfyx, { 2, 2, 8, 3 }, format::bfyx, { 2, 2, 8, 3 }, 0, data_types::f32, format::bfyx +#define CASE_GATHER_ELEMENTS_FP32_4D_3 data_types::f32, format::bfyx, { 1, 3, 2, 9 }, format::bfyx, { 1, 3, 5, 9 }, format::bfyx, { 1, 3, 5, 9 }, 3, data_types::f32, format::bfyx -#define CASE_GATHER_ELEMENTS_FP32_5D_1 data_types::f32, format::bfzyx, { 3, 2, 5, 2, 3 }, format::bfzyx, { 3, 2, 2, 2, 3 }, format::bfzyx, { 3, 2, 2, 2, 3 }, cldnn::gather_elements::gather_elements_axis::along_x, data_types::f32, format::bfzyx -#define CASE_GATHER_ELEMENTS_FP32_5D_2 data_types::f32, format::bfzyx, { 5, 4, 7, 4, 4 }, format::bfzyx, { 5, 4, 7, 4, 3 }, format::bfzyx, { 5, 4, 7, 4, 3 }, cldnn::gather_elements::gather_elements_axis::along_z, data_types::f32, format::bfzyx +#define CASE_GATHER_ELEMENTS_FP32_5D_1 data_types::f32, format::bfzyx, { 3, 2, 5, 2, 3 }, format::bfzyx, { 3, 2, 2, 2, 3 }, format::bfzyx, { 3, 2, 2, 2, 3 }, 4, data_types::f32, format::bfzyx +#define CASE_GATHER_ELEMENTS_FP32_5D_2 data_types::f32, format::bfzyx, { 5, 4, 7, 4, 4 }, format::bfzyx, { 5, 4, 7, 4, 3 }, format::bfzyx, { 5, 4, 7, 4, 3 }, 2, data_types::f32, format::bfzyx -#define CASE_GATHER_ELEMENTS_FP32_6D_1 data_types::f32, format::bfwzyx, { 5, 4, 6, 7, 8, 2 }, format::bfwzyx, { 5, 2, 6, 7, 8, 2 }, format::bfwzyx, { 5, 2, 6, 7, 8, 2 }, cldnn::gather_elements::gather_elements_axis::along_f, data_types::f32, format::bfwzyx -#define CASE_GATHER_ELEMENTS_FP32_6D_2 data_types::f32, format::bfwzyx, { 2, 1, 2, 3, 2, 1 }, format::bfwzyx, { 2, 1, 2, 3, 2, 3 }, format::bfwzyx, { 2, 1, 2, 3, 2, 3 }, cldnn::gather_elements::gather_elements_axis::along_w, data_types::f32, format::bfwzyx -#define CASE_GATHER_ELEMENTS_FP32_6D_3 data_types::f32, format::bfwzyx, { 2, 2, 3, 4, 4, 2 }, format::bfwzyx, { 2, 2, 6, 4, 4, 2 }, format::bfwzyx, { 2, 2, 6, 4, 4, 2 }, cldnn::gather_elements::gather_elements_axis::along_x, data_types::f32, format::bfwzyx +#define CASE_GATHER_ELEMENTS_FP32_6D_1 data_types::f32, format::bfwzyx, { 5, 4, 6, 7, 8, 2 }, format::bfwzyx, { 5, 2, 6, 7, 8, 2 }, format::bfwzyx, { 5, 2, 6, 7, 8, 2 }, 1, data_types::f32, format::bfwzyx +#define CASE_GATHER_ELEMENTS_FP32_6D_2 data_types::f32, format::bfwzyx, { 2, 1, 2, 3, 2, 1 }, format::bfwzyx, { 2, 1, 2, 3, 2, 3 }, format::bfwzyx, { 2, 1, 2, 3, 2, 3 }, 2, data_types::f32, format::bfwzyx +#define CASE_GATHER_ELEMENTS_FP32_6D_3 data_types::f32, format::bfwzyx, { 2, 2, 3, 4, 4, 2 }, format::bfwzyx, { 2, 2, 6, 4, 4, 2 }, format::bfwzyx, { 2, 2, 6, 4, 4, 2 }, 5, data_types::f32, format::bfwzyx class gather_elements_quantize : public GatherElementsPrimitiveFusingTest {}; TEST_P(gather_elements_quantize, basic) { diff --git a/src/plugins/intel_gpu/tests/test_cases/gather_elements_gpu_test.cpp b/src/plugins/intel_gpu/tests/test_cases/gather_elements_gpu_test.cpp index dee59f8a978..6b531043b07 100644 --- a/src/plugins/intel_gpu/tests/test_cases/gather_elements_gpu_test.cpp +++ b/src/plugins/intel_gpu/tests/test_cases/gather_elements_gpu_test.cpp @@ -21,7 +21,7 @@ inline void DoTest(engine& engine, const cldnn::memory::ptr& input1, // indices const std::vector& expected_results, const tensor& output_tensor, - const cldnn::gather_elements::gather_elements_axis axis) { + const int64_t axis) { topology topology; topology.add(input_layout("InputData", input0->get_layout())); topology.add(input_layout("InputIndices", input1->get_layout())); @@ -45,7 +45,7 @@ inline void DoTest(engine& engine, TEST(gather_elements_gpu_fp16, d3283_i2283_a0) { auto& engine = get_test_engine(); - auto axis = cldnn::gather_elements::gather_elements_axis::along_b; + auto axis = 0; auto input0 = engine.allocate_memory({ data_types::f16, format::bfyx, { 3, 2, 8, 3 } }); // data auto input1 = engine.allocate_memory({ data_types::f16, format::bfyx, { 2, 2, 8, 3 } }); // indices @@ -106,7 +106,7 @@ TEST(gather_elements_gpu_fp16, d3283_i2283_a0) { TEST(gather_elements_gpu_fp16, d2235_i2235_a3) { auto& engine = get_test_engine(); - auto axis = cldnn::gather_elements::gather_elements_axis::along_x; + auto axis = 3; auto input0 = engine.allocate_memory({ data_types::f16, format::bfyx, { 2, 2, 3, 5 } }); // data auto input1 = engine.allocate_memory({ data_types::f16, format::bfyx, { 2, 2, 3, 5 } }); // indices set_values(input0, { @@ -184,7 +184,7 @@ TEST(gather_elements_gpu_fp16, d2235_i2235_a3) { TEST(gather_elements_gpu_fp16, d1329_i1359_an1) { auto& engine = get_test_engine(); - auto axis = cldnn::gather_elements::gather_elements_axis::along_x; + auto axis = 3; auto input0 = engine.allocate_memory({ data_types::f16, format::bfyx, { 1, 3, 2, 9 } }); // data auto input1 = engine.allocate_memory({ data_types::f16, format::bfyx, { 1, 3, 5, 9 } }); // indices set_values(input0, { @@ -283,7 +283,7 @@ TEST(gather_elements_gpu_fp16, d1329_i1359_an1) { TEST(gather_elements_gpu_fp16, d12853_i12923_a3) { auto& engine = get_test_engine(); - auto axis = cldnn::gather_elements::gather_elements_axis::along_y; + auto axis = 3; auto input0 = engine.allocate_memory({ data_types::f16, format::bfzyx, { 1, 2, 8, 5, 3 } }); // data auto input1 = engine.allocate_memory({ data_types::f16, format::bfzyx, { 1, 2, 8, 2, 3 } }); // indices @@ -356,7 +356,7 @@ TEST(gather_elements_gpu_fp16, d12853_i12923_a3) { TEST(gather_elements_gpu_fp16, d25441_i22441_an4) { auto& engine = get_test_engine(); - auto axis = cldnn::gather_elements::gather_elements_axis::along_f; + auto axis = 1; auto input0 = engine.allocate_memory({ data_types::f16, format::bfzyx, { 2, 5, 4, 4, 1 } }); // data auto input1 = engine.allocate_memory({ data_types::f16, format::bfzyx, { 2, 2, 4, 4, 1 } }); // indices @@ -448,7 +448,7 @@ TEST(gather_elements_gpu_fp16, d25441_i22441_an4) { TEST(gather_elements_gpu_fp16, d32843_i12843_a0) { auto& engine = get_test_engine(); - auto axis = cldnn::gather_elements::gather_elements_axis::along_b; + auto axis = 0; auto input0 = engine.allocate_memory({ data_types::f16, format::bfzyx, { 3, 2, 8, 4, 3 } }); // data auto input1 = engine.allocate_memory({ data_types::f16, format::bfzyx, { 1, 2, 8, 4, 3 } }); // indices @@ -588,7 +588,7 @@ TEST(gather_elements_gpu_fp16, d32843_i12843_a0) { TEST(gather_elements_gpu_fp16, d223442_i226442_a5) { auto& engine = get_test_engine(); - auto axis = cldnn::gather_elements::gather_elements_axis::along_x; + auto axis = 5; auto input0 = engine.allocate_memory({ data_types::f16, format::bfwzyx, { 2, 2, 3, 4, 4, 2 } }); // data auto input1 = engine.allocate_memory({ data_types::f16, format::bfwzyx, { 2, 2, 6, 4, 4, 2 } }); // indices @@ -991,7 +991,7 @@ TEST(gather_elements_gpu_fp16, d223442_i226442_a5) { TEST(gather_elements_gpu_fp16, d124251_i124221_an3) { auto& engine = get_test_engine(); - auto axis = cldnn::gather_elements::gather_elements_axis::along_z; + auto axis = 3; auto input0 = engine.allocate_memory({ data_types::f16, format::bfwzyx, { 1, 2, 4, 2, 5, 1 } }); // data auto input1 = engine.allocate_memory({ data_types::f16, format::bfwzyx, { 1, 2, 4, 2, 2, 1 } }); // indices @@ -1046,7 +1046,7 @@ TEST(gather_elements_gpu_fp16, d124251_i124221_an3) { TEST(gather_elements_gpu_fp16, d233113_i233115_a2) { auto& engine = get_test_engine(); - auto axis = cldnn::gather_elements::gather_elements_axis::along_w; + auto axis = 2; auto input0 = engine.allocate_memory({ data_types::f16, format::bfwzyx, { 2, 3, 3, 1, 1, 3 } }); // data auto input1 = engine.allocate_memory({ data_types::f16, format::bfwzyx, { 2, 3, 3, 1, 1, 5 } }); // indices