Skip to content

Fix LSTM conversion for models with rank > 3 inputs from Unsqueeze operations #33023

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 7 commits into
base: master
Choose a base branch
from
+1,108 −3
Open

Fix LSTM conversion for models with rank > 3 inputs from Unsqueeze operations #33023

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
d1689ce
Add rank reduction for LSTM inputs to enable silero_vad.onnx conversion
evkotov Nov 21, 2025
de3353d
add lstm test
evkotov Nov 25, 2025
b06e7f2
remove unneeded lstm_rank5_reshape test
evkotov Nov 25, 2025
c58c5bc
add missing lstm_rank4_with_unsqueeze.prototxt
evkotov Nov 25, 2025
5ba2f09
clang fix
evkotov Nov 25, 2025
08c7dda
code review fix: remove unneeded argument
evkotov Nov 27, 2025
5bc7dca
code review fix: use Output<Node> return type, always use Squeeze
evkotov Dec 3, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 48 additions & 3 deletions src/frontends/onnx/frontend/src/op/lstm.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,8 @@
#include "openvino/op/lstm_sequence.hpp"
#include "openvino/op/multiply.hpp"
#include "openvino/op/shape_of.hpp"
#include "openvino/op/squeeze.hpp"
#include "openvino/op/tile.hpp"
#include "openvino/util/common_util.hpp"
#include "utils/reshape.hpp"
#include "utils/split.hpp"
Expand All @@ -37,6 +39,32 @@ enum class LSTMInput {
LSTM_INPUT_P
};

// Helper function to reduce tensor rank to target_rank by squeezing leading dimensions.
// Per ONNX specification, LSTM requires rank-3 inputs, so extra dimensions must be == 1.
// If extra dimensions are != 1 at runtime, Squeeze will fail with a clear error.
ov::Output<ov::Node> reduce_tensor_rank(const ov::Output<ov::Node>& input, int64_t target_rank) {
const auto& input_shape = input.get_partial_shape();

if (input_shape.rank().is_dynamic()) {
return input;
}

const auto input_rank = input_shape.rank().get_length();

if (input_rank <= target_rank) {
return input;
}

// Squeeze all leading dimensions to reduce rank to target_rank
std::vector<int64_t> axes_to_squeeze;
for (int64_t i = 0; i < input_rank - target_rank; ++i) {
axes_to_squeeze.push_back(i);
}

auto axes_const = v0::Constant::create(ov::element::i64, Shape{axes_to_squeeze.size()}, axes_to_squeeze);
return std::make_shared<v0::Squeeze>(input, axes_const);
}

struct LSTMNgInputMap {
explicit LSTMNgInputMap(const Node& node) {
const auto& ng_inputs = node.get_ov_inputs();
Expand All @@ -48,7 +76,14 @@ struct LSTMNgInputMap {
// Packed input sequences.
// ONNX Shape: [seq_length, batch_size, input_size]
// OpenVino Shape: [batch_size, seq_length, input_size]
m_input_map[LSTMInput::LSTM_INPUT_X] = ov::op::util::reorder_axes(ng_inputs.at(0), {1, 0, 2});

// First reduce rank if needed, THEN reorder axes
// This is important because Squeeze changes dimension indices
auto input_x = ng_inputs.at(0);
input_x = reduce_tensor_rank(input_x, 3);
input_x = ov::op::util::reorder_axes(input_x, {1, 0, 2});

m_input_map[LSTMInput::LSTM_INPUT_X] = input_x;

// Weight tensor for the gates.
// Shape: [num_directions, 4*hidden_size, input_size]
Expand Down Expand Up @@ -124,7 +159,12 @@ struct LSTMNgInputMap {
// ONNX Shape: [num_directions, batch_size, hidden_size]
// OpenVino Shape: [batch_size, num_directions, hidden_size]
if (ng_inputs.size() > 5 && !ov::op::util::is_null(ng_inputs.at(5))) {
m_input_map[LSTMInput::LSTM_INPUT_INIT_H] = ov::op::util::reorder_axes(ng_inputs.at(5), {1, 0, 2});
auto init_h = ng_inputs.at(5);
// First reduce rank, THEN reorder axes
init_h = reduce_tensor_rank(init_h, 3);
init_h = ov::op::util::reorder_axes(init_h, {1, 0, 2});

m_input_map[LSTMInput::LSTM_INPUT_INIT_H] = init_h;
} else {
auto init_h_shape =
std::make_shared<v0::Concat>(ov::OutputVector{batch_size_node, num_directions_node, hidden_size_node},
Expand All @@ -137,7 +177,12 @@ struct LSTMNgInputMap {
// ONNX Shape: [num_directions, batch_size, hidden_size]
// OpenVino Shape: [batch_size, num_directions, hidden_size]
if (ng_inputs.size() > 6 && !ov::op::util::is_null(ng_inputs.at(6))) {
m_input_map[LSTMInput::LSTM_INPUT_INIT_C] = ov::op::util::reorder_axes(ng_inputs.at(6), {1, 0, 2});
auto init_c = ng_inputs.at(6);
// First reduce rank, THEN reorder axes
init_c = reduce_tensor_rank(init_c, 3);
init_c = ov::op::util::reorder_axes(init_c, {1, 0, 2});

m_input_map[LSTMInput::LSTM_INPUT_INIT_C] = init_c;
} else {
auto init_c_shape =
std::make_shared<v0::Concat>(ov::OutputVector{batch_size_node, num_directions_node, hidden_size_node},
Expand Down
167 changes: 167 additions & 0 deletions src/frontends/onnx/tests/models/lstm_high_rank_input.prototxt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,167 @@
ir_version: 4
producer_name: "OpenVINO ONNX Frontend"
graph {
node {
input: "X"
input: "W"
input: "R"
input: "B"
output: "Y"
output: "Y_h"
output: "Y_c"
op_type: "LSTM"
attribute {
name: "direction"
s: "forward"
type: STRING
}
attribute {
name: "hidden_size"
i: 2
type: INT
}
}
name: "compute_graph"
input {
name: "X"
type {
tensor_type {
elem_type: 1
shape {
dim {
dim_value: 1
}
dim {
dim_value: 1
}
dim {
dim_value: 3
}
dim {
dim_value: 2
}
dim {
dim_value: 4
}
}
}
}
}
input {
name: "W"
type {
tensor_type {
elem_type: 1
shape {
dim {
dim_value: 1
}
dim {
dim_value: 8
}
dim {
dim_value: 4
}
}
}
}
}
input {
name: "R"
type {
tensor_type {
elem_type: 1
shape {
dim {
dim_value: 1
}
dim {
dim_value: 8
}
dim {
dim_value: 2
}
}
}
}
}
input {
name: "B"
type {
tensor_type {
elem_type: 1
shape {
dim {
dim_value: 1
}
dim {
dim_value: 16
}
}
}
}
}
output {
name: "Y"
type {
tensor_type {
elem_type: 1
shape {
dim {
dim_value: 3
}
dim {
dim_value: 1
}
dim {
dim_value: 2
}
dim {
dim_value: 2
}
}
}
}
}
output {
name: "Y_h"
type {
tensor_type {
elem_type: 1
shape {
dim {
dim_value: 1
}
dim {
dim_value: 2
}
dim {
dim_value: 2
}
}
}
}
}
output {
name: "Y_c"
type {
tensor_type {
elem_type: 1
shape {
dim {
dim_value: 1
}
dim {
dim_value: 2
}
dim {
dim_value: 2
}
}
}
}
}
}
opset_import {
version: 7
}
Loading
Loading