fix test and op spec

Author: zhangYiIntel

src/bindings/python/src/openvino/_pyopenvino/op/__init__.pyi

@@ -182,8 +182,7 @@ class _PagedAttentionExtension(openvino._pyopenvino.Node):
     Experimental extention for PagedAttention operation. Use with care: no backward compatibility is guaranteed in future releases.
     """
     def __init__(self, arg0: collections.abc.Sequence[openvino._pyopenvino.Output]) -> None:
-        ...
-        
+        ...        
 class _GatedDeltaNet(openvino._pyopenvino.Node):
     """
     Experimental extention for GatedDeltaNet operation. Use with care: no backward compatibility is guaranteed in future releases.

src/core/src/op/gated_delta_net.cpp

@@ -77,10 +77,71 @@ void GatedDeltaNet::validate_and_infer_types() {
     input_check(this, 4, "gate", {3}, {});
     input_check(this, 5, "beta", {3}, {});
 
-    // value head_size may be not same with key
+    // batch, seq_len, head_num, head_size
+    const auto& query_ps = get_input_partial_shape(0);
+    const auto& key_ps = get_input_partial_shape(1);
+    const auto& value_ps = get_input_partial_shape(2);
+    const auto& state_ps = get_input_partial_shape(3);
+    const auto& gate_ps = get_input_partial_shape(4);
+    const auto& beta_ps = get_input_partial_shape(5);
+    
+    const auto q_head_num = query_ps[2];
+    const auto k_head_num = key_ps[2];
+    const auto v_head_num = value_ps[2];
+
+    const auto k_head_size = key_ps[3];
+    const auto v_head_size = value_ps[3];
+
+    NODE_VALIDATION_CHECK(this, q_head_num.is_static() && k_head_num.is_static() && q_head_num.get_length() == k_head_num.get_length(),
+                          "The number of heads in query and key should be the same, but got ",
+                          q_head_num,
+                          " and ",
+                          k_head_num,
+                          ".");
+
+    NODE_VALIDATION_CHECK(this, k_head_size.is_static() && v_head_size.is_static() && k_head_size.get_length() == v_head_size.get_length(),
+                          "The head size in key and value should be the same, but got ",
+                          k_head_size,
+                          " and ",
+                          v_head_size,
+                          ".");
+
+    const auto gate_head_num = gate_ps[2];
+    const auto beta_head_num = beta_ps[2];
+
+    NODE_VALIDATION_CHECK(this, gate_head_num.is_static() && beta_head_num.is_static() && gate_head_num.get_length() == beta_head_num.get_length(),
+                          "The number of heads in gate and beta should be the same, but got ",
+                          gate_head_num,
+                          " and ",
+                          beta_head_num,
+                          ".");
+
+    // [batch, v_head_nums, v_head_size, k_head_size]
+    const auto state_head_num = state_ps[1];
+    const auto state_hidden_size_0 = state_ps[2];
+    const auto state_hidden_size_1 = state_ps[3];
+    NODE_VALIDATION_CHECK(this, state_head_num.is_static() && state_head_num.get_length() == v_head_num.get_length(),
+                          "The number of heads in recurrent_state and value should be the same, but got ",
+                          state_head_num,
+                          " and ",
+                          v_head_num,
+                          ".");
+    NODE_VALIDATION_CHECK(this, state_hidden_size_0.is_static() && state_hidden_size_0.get_length() == v_head_size.get_length(),
+                          "The [-2] dim in shape of recurrent_state and value should be the same, but got ",
+                          state_hidden_size_0,
+                          " and ",
+                          v_head_size,
+                          ".");
+    NODE_VALIDATION_CHECK(this, state_hidden_size_1.is_static() && state_hidden_size_1.get_length() == k_head_size.get_length(),
+                          "The [-1] dim in shape of recurrent_state and key should be the same, but got ",
+                          state_hidden_size_1,
+                          " and ",
+                          k_head_size,
+                          ".");
+    // output has the same shape and type as input value, output state has the same shape and type as input recurrent_state
     auto out_ps = get_input_partial_shape(2);
     const auto& h_ps = get_input_partial_shape(3);
-    set_output_type(0, get_input_element_type(0), out_ps);
+    set_output_type(0, get_input_element_type(2), out_ps);
     set_output_type(1, get_input_element_type(3), h_ps);
 }
 

src/tests/functional/plugin/shared/include/subgraph_tests/gated_delta_net.hpp

@@ -9,10 +9,26 @@
 namespace ov {
 namespace test {
 
+inline void CheckNumberOfNodesWithType(std::shared_ptr<const ov::Model> function,
+                                       const std::unordered_set<std::string>& nodeTypes,
+                                       size_t expectedCount) {
+    ASSERT_NE(nullptr, function);
+    int num_ops = 0;
+    for (const auto& node : function->get_ordered_ops()) {
+        const auto& rt_info = node->get_rt_info();
+        const auto layer_type = rt_info.find("layerType")->second.as<std::string>();
+        if (nodeTypes.count(layer_type)) {
+            num_ops++;
+        }
+    }
+    ASSERT_EQ(num_ops, expectedCount);
+}
+
 TEST_P(GatedDeltaNet, CompareWithRefs) {
     SKIP_IF_CURRENT_TEST_IS_DISABLED();
     run();
     auto function = compiledModel.get_runtime_model();
+    CheckNumberOfNodesWithType(function, {"GatedDeltaNet"}, 1);
 };
 }  // namespace test
 }  // namespace ov