[Layer][GPU][TF FE] test_segment_sum failed with get_shape was called…

… on a descriptor

Fix GPU TensorFlow SegmentSum test failure
- Added dynamic shape support in GPU inference pipeline for SegmentSum ops
- Updated tensorflow_tests/test_tf_SegmentSum.py to remove GPU skip
- Ensured correct handling of dynamic shapes during GPU inference for SegmentSum
This update resolves the GPU test failure and ensures proper support for dynamic shapes in SegmentSum ops.

Jira Ticket:
https://jira.devtools.intel.com/browse/CVS-105896

src/core/shape_inference/include/embedding_segments_sum_shape_inference.hpp

@@ -53,7 +53,12 @@ std::vector<TRShape> shape_infer(const EmbeddingSegmentsSum* op,
     if (emb_table_shape.rank().is_static()) {
         NODE_VALIDATION_CHECK(op, emb_table_shape.size() > 0, "EMB_TABLE can't be a scalar.");
         if (auto segments_value = get_input_const_data_as_shape<TRShape>(op, NUM_SEGMENTS, ta)) {
-            result_shape[0] = (*segments_value)[0];
+            if (segments_value->size() > 1) {
+                size_t sz = segments_value->size() - 1;
+                result_shape[0] = (*segments_value)[sz] + 1;
+            } else {
+                result_shape[0] = (*segments_value)[0];
+            }
         } else {
             result_shape[0] = Dimension::dynamic();
         }

src/plugins/intel_gpu/src/graph/embedding_bag.cpp

@@ -29,6 +29,7 @@ std::vector<layout> embedding_bag_inst::calc_output_layouts(embedding_bag_node c
     for (size_t i = 0; i < desc->input_size(); i++) {
         input_shapes.push_back(impl_param.get_input_layout(i).get<ShapeType>());
     }
+    input_shapes.push_back(ShapeType());
 
     std::vector<ShapeType> output_shapes;
 
@@ -46,7 +47,7 @@ std::vector<layout> embedding_bag_inst::calc_output_layouts(embedding_bag_node c
         case embedding_bag::segments_sum: {
             ov::op::v3::EmbeddingSegmentsSum op;
 
-            const size_t num_segments_idx = 3;
+            const size_t num_segments_idx = 2;
             TensorsContainer const_data(&impl_param.get_stream());
             if (memory_deps.count(num_segments_idx) > 0) {
                 const_data.emplace(3, memory_deps.at(num_segments_idx));

src/plugins/intel_gpu/src/graph/impls/ocl/embedding_bag.cpp

@@ -56,6 +56,7 @@ attach_embedding_bag_impl::attach_embedding_bag_impl() {
     implementation_map<embedding_bag>::add(impl_types::ocl, typed_primitive_impl_ocl<embedding_bag>::create<embedding_bag_impl>, {
         std::make_tuple(data_types::f32, format::bfyx),
         std::make_tuple(data_types::f16, format::bfyx),
+        std::make_tuple(data_types::i32, format::bfyx),
     });
 }
 

src/plugins/intel_gpu/src/kernel_selector/kernels/embedding_bag/embedding_bag_kernel_ref.cpp

@@ -86,8 +86,10 @@ ParamsKey EmbeddingBagKernelRef::GetSupportedKey() const {
     k.EnableInputDataType(Datatype::INT32);
     k.EnableInputDataType(Datatype::INT64);
     k.EnableInputDataType(Datatype::UINT32);
+
     k.EnableOutputDataType(Datatype::F16);
     k.EnableOutputDataType(Datatype::F32);
+    k.EnableOutputDataType(Datatype::INT32);
 
     k.EnableAllInputLayout();
     k.EnableAllOutputLayout();

src/plugins/intel_gpu/src/plugin/ops/embedding_bag.cpp

@@ -141,10 +141,13 @@ static void CreateEmbeddingSegmentsSumOp(ProgramBuilder& p, const std::shared_pt
         }
     }
 
+    auto p_shape = op->get_output_partial_shape(0);
+    auto output_shape = p_shape.is_static() ? tensor_from_dims(p_shape.to_shape()) : cldnn::tensor();
+
     auto embeddingBagPrim = cldnn::embedding_bag(layerName,
                                                  reordered_inputs,
                                                  cldnn::embedding_bag::segments_sum,
-                                                 tensor_from_dims(op->get_output_shape(0)),
+                                                 output_shape,
                                                  defaultIndex);
 
     p.add_primitive(*op, embeddingBagPrim);

tests/layer_tests/tensorflow_tests/test_tf_SegmentSum.py

@@ -49,8 +49,6 @@ def test_segment_sum_basic(self, params, ie_device, precision, ir_version, temp_
                                use_legacy_frontend):
         if use_legacy_frontend:
             pytest.skip("SegmentSum operation is not supported via legacy frontend.")
-        if ie_device == 'GPU':
-            pytest.skip("GPU error: to_shape was called on a dynamic shape")
         self._test(*self.create_segment_sum_net(**params),
                    ie_device, precision, ir_version, temp_dir=temp_dir,
                    use_legacy_frontend=use_legacy_frontend)
@@ -68,8 +66,6 @@ def test_segment_sum_different_types(self, params, ie_device, precision, ir_vers
                                          use_legacy_frontend):
         if use_legacy_frontend:
             pytest.skip("SegmentSum operation is not supported via legacy frontend.")
-        if ie_device == 'GPU':
-            pytest.skip("GPU error: to_shape was called on a dynamic shape")
         self._test(*self.create_segment_sum_net(**params),
                    ie_device, precision, ir_version, temp_dir=temp_dir,
                    use_legacy_frontend=use_legacy_frontend)
@@ -122,8 +118,6 @@ def test_complex_segment_sum(self, params, ie_device, precision, ir_version, tem
                                          use_legacy_frontend):
         if use_legacy_frontend:
             pytest.skip("SegmentSum operation is not supported via legacy frontend.")
-        if ie_device == 'GPU':
-            pytest.skip("GPU error: to_shape was called on a dynamic shape")
         self._test(*self.create_segment_sum_net(**params),
                    ie_device, precision, ir_version, temp_dir=temp_dir,
-                   use_legacy_frontend=use_legacy_frontend)
+                   use_legacy_frontend=use_legacy_frontend)