Stateful fix for shape errors after rebase

Author: cavusmustafa

ggml/src/ggml-openvino/openvino/op/glu_swiglu.cpp

@@ -26,7 +26,12 @@ OutputVector translate_glu_swiglu(const NodeContext & context) {
         src1 = context.get_input(1);
     } else {
         auto combined = context.get_input(0);
-        auto split_axis = ov::op::v0::Constant::create(ov::element::i64, {}, {3});
+       // TODO: Will it work if we set it to "-1" for all cases?
+        int32_t split_dim = 3;
+        if (context.is_stateful()) {
+            split_dim = -1;
+        }
+        auto split_axis = ov::op::v0::Constant::create(ov::element::i64, {}, {split_dim});
         auto split = std::make_shared<ov::op::v1::Split>(combined, split_axis, 2);
         src0 = split->output(0);
         src1 = split->output(1);

ggml/src/ggml-openvino/openvino/op/rope.cpp

@@ -104,8 +104,12 @@ OutputVector translate_rope(const NodeContext & context) {
             ov::element::i64, {4}, std::vector<int64_t>{1, -1, (int64_t) output_shape[2], (int64_t) output_shape[3]});
         res = std::make_shared<ov::op::v1::Reshape>(stack, data_shape, false);
     } else if (mode == ROPE_TYPE_NEOX) {
+        int32_t split_dim = 3;
+        if (context.is_stateful()) {
+            split_dim = 2;
+        }
         auto data_split = std::make_shared<ov::op::v1::Split>(
-            data_node, ov::op::v0::Constant::create(ov::element::i64, ov::Shape{}, {3}), 2);
+            data_node, ov::op::v0::Constant::create(ov::element::i64, ov::Shape{}, {split_dim}), 2);
         Output<Node> slice_data_node_0 = data_split->outputs()[0];
         Output<Node> slice_data_node_1 = data_split->outputs()[1];
 
@@ -117,9 +121,10 @@ OutputVector translate_rope(const NodeContext & context) {
             std::make_shared<ov::op::v1::Multiply>(slice_data_node_0, sin_theta_node),
             std::make_shared<ov::op::v1::Multiply>(slice_data_node_1, cos_theta_node));
 
+        // TODO: Will it work if we set it to "-1" for all cases?
         int32_t concat_dim = 3;
         if (context.is_stateful()) {
-            concat_dim = 2;
+            concat_dim = -1;
         }
         res = std::make_shared<ov::op::v0::Concat>(ov::OutputVector{first_half_node, second_half_node}, concat_dim);
     }

ggml/src/ggml-openvino/openvino/utils.cpp

@@ -213,10 +213,15 @@ ov::Output<ov::Node> process_view_input(const NodeContext & context, int input_i
     }
     int64_t slice_end = split_addr + slice_len;
 
+    int32_t axes_val = 3;
+    if (context.is_stateful()) {
+        axes_val = 2;
+    }
+
     auto begin = ov::op::v0::Constant::create(ov::element::i64, {1}, {split_addr});
     auto end = ov::op::v0::Constant::create(ov::element::i64, {1}, {slice_end});
     auto stride = ov::op::v0::Constant::create(ov::element::i64, {1}, {1});
-    auto axes = ov::op::v0::Constant::create(ov::element::i64, {1}, {3});
+    auto axes = ov::op::v0::Constant::create(ov::element::i64, {1}, {axes_val});
     auto sliced = std::make_shared<ov::op::v8::Slice>(input, begin, end, stride, axes);
     return sliced;
 }

ggml/src/ggml-openvino/utils.cpp

@@ -614,10 +614,10 @@ ov::Tensor get_ov_output_tensor(std::shared_ptr<GgmlOvDecoder> ggml_decoder, con
     auto output_type = ggml_decoder->get_ov_type(ggml_tensor);
     auto output_shape = ggml_decoder->get_shape(ggml_tensor);
 
-    if (ggml_decoder->is_static() && result_name == "result_output" && output_shape[2] == 0) {
+    if (ggml_decoder->is_static() && (result_name == "result_output" || result_name == "result_norm") && output_shape[2] == 0) {
         output_shape[2] = 1;
     }
-    if (ggml_decoder->is_stateful() && result_name == "result_output") {
+    if (ggml_decoder->is_stateful() && (result_name == "result_output" || result_name == "result_norm")) {
         std::vector<long unsigned int> output_shape_3d;
         for (size_t i=1; i<output_shape.size(); i++) {
             output_shape_3d.push_back(output_shape[i]);