Skip to content

[OpenVINO backend] Support for LSTM #22313

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

Closed
andersendsa wants to merge 16 commits into from
Closed

[OpenVINO backend] Support for LSTM #22313

Show file tree
Hide file tree
Changes from 14 commits
Commits
Show all changes
16 commits
Select commit Hold shift + click to select a range
5f9c2df
Implement lstm in OpenVINO Backend
andersendsa Feb 28, 2026
ac9384e
fixing format
andersendsa Feb 28, 2026
d08e3db
unsqueeze final output
andersendsa Feb 28, 2026
d8b91b8
resolve circular dependancy
andersendsa Feb 28, 2026
9683ce2
resolve circular dependancy
andersendsa Feb 28, 2026
f4ab9bc
resolve circular dependancy
andersendsa Feb 28, 2026
c2e50b1
formatting fix
andersendsa Feb 28, 2026
107f5a8
formatting fix
andersendsa Feb 28, 2026
c6cff76
resolving unnessary changes
andersendsa Mar 4, 2026
4e5eee4
reverting changes
andersendsa Mar 7, 2026
8c2fd6a
excluding tests
andersendsa Mar 8, 2026
15d8b3b
Merge branch 'master' into openvino-lstm
andersendsa Mar 8, 2026
7943a2c
excluding tests
andersendsa Mar 8, 2026
5a2f450
fixing masking tolerances
andersendsa Mar 8, 2026
261acd8
fixing masking tolerances
andersendsa Mar 8, 2026
486ecb6
fixing error
andersendsa Mar 8, 2026
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
5 changes: 0 additions & 5 deletions keras/src/backend/openvino/excluded_concrete_tests.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -421,11 +421,6 @@ LossScaleOptimizerTest::test_iterations_update_stateful
LossScaleOptimizerTest::test_iterations_update_stateless
LossScaleOptimizerTest::test_upscaling_stateful
LossScaleOptimizerTest::test_upscaling_stateless
LSTMTest::test_correctness0
LSTMTest::test_correctness1
LSTMTest::test_masking
LSTMTest::test_pass_initial_state
LSTMTest::test_statefulness
MathOpsCorrectnessTest::test_erfinv_operation_basic
MathOpsCorrectnessTest::test_erfinv_operation_dtype
MathOpsCorrectnessTest::test_logdet
Expand Down
190 changes: 180 additions & 10 deletions keras/src/backend/openvino/rnn.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
from openvino import Model
from openvino import Type

from keras.src import activations
from keras.src import tree
from keras.src.backend.openvino.core import OpenVINOKerasTensor
from keras.src.backend.openvino.core import get_ov_output
Expand Down Expand Up @@ -105,7 +106,7 @@ def _slice_at_0(x):
x_ov = get_ov_output(x)
slice_0 = ov_opset.gather(
x_ov,
ov_opset.constant([0], dtype=Type.i32).output(0),
ov_opset.constant(0, dtype=Type.i32).output(0),
ov_opset.constant(0, dtype=Type.i32).output(0),
).output(0)
return OpenVINOKerasTensor(slice_0)
Expand All @@ -130,10 +131,10 @@ def _slice_at_0(x):
inp_ov = get_ov_output(inp)
pshape = inp_ov.get_partial_shape()
if pshape.rank.is_static:
new_shape = list(pshape)[1:]
new_shape = [1] + list(pshape)[1:]
else:
new_shape = (
[-1] * (pshape.rank.get_length() - 1)
[-1] * (pshape.rank.get_length())
if pshape.rank.is_static
else None
)
Expand All @@ -144,7 +145,7 @@ def _slice_at_0(x):
if mask is not None:
mask_ov = get_ov_output(mask)
pshape = mask_ov.get_partial_shape()
new_shape = list(pshape)[1:] if pshape.rank.is_static else None
new_shape = [1] + list(pshape)[1:] if pshape.rank.is_static else None
param = ov_opset.parameter(new_shape, mask_ov.get_element_type())
sliced_mask_params.append(param)
params.append(param)
Expand Down Expand Up @@ -172,9 +173,13 @@ def _slice_at_0(x):
)
constants_params.append(param)
params.append(param)
sliced_inputs_t = [
OpenVINOKerasTensor(p.output(0)) for p in sliced_inputs_params
]
sliced_inputs_t = []
for p in sliced_inputs_params:
p_out = p.output(0)
p_out = ov_opset.squeeze(
p_out, ov_opset.constant([0], dtype=Type.i32).output(0)
).output(0)
sliced_inputs_t.append(OpenVINOKerasTensor(p_out))
merged_states_t = [
OpenVINOKerasTensor(p.output(0)) for p in merged_states_params
]
Expand All @@ -192,6 +197,9 @@ def _slice_at_0(x):
final_last_output_list = []
if mask is not None:
mask_t = sliced_mask_params[0].output(0)
mask_t = ov_opset.squeeze(
mask_t, ov_opset.constant([0], dtype=Type.i32).output(0)
).output(0)
for i, (new_st, old_st) in enumerate(
zip(flat_step_new_states, merged_states_t)
):
Expand Down Expand Up @@ -225,7 +233,15 @@ def _slice_at_0(x):
final_states_list = [get_ov_output(x) for x in flat_step_new_states]
final_output_list = [get_ov_output(x) for x in flat_step_output]
final_last_output_list = [get_ov_output(x) for x in flat_step_output]
zero_const = ov_opset.constant([0], dtype=Type.i32).output(0)
final_output_list = [
ov_opset.unsqueeze(x, zero_const).output(0) for x in final_output_list
]
cond_const = ov_opset.constant(True, Type.boolean).output(0)
zero_const = ov_opset.constant([0], dtype=Type.i32).output(0)
final_output_list = [
ov_opset.unsqueeze(x, zero_const).output(0) for x in final_output_list
]
Comment on lines +241 to +244
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot Mar 8, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

This block of code, which redefines zero_const and re-assigns final_output_list by unsqueezing its elements, is a duplicate of the block on lines 236-239. This is redundant and should be removed to avoid confusion and potential bugs.

results = (
[cond_const]
+ final_states_list
Expand Down Expand Up @@ -284,8 +300,162 @@ def _slice_at_0(x):
return last_output, outputs, new_states


def lstm(*args, **kwargs):
raise NotImplementedError("`lstm` is not supported with openvino backend")
def lstm(
inputs,
initial_state_h,
initial_state_c,
mask,
kernel,
recurrent_kernel,
bias,
activation,
recurrent_activation,
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot Feb 28, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The unroll parameter is included in the function signature but is not used within the implementation. According to the Keras API design guidelines (lines 18 and 53), unused parameters should be avoided to keep the API surface clean and minimize user confusion. Please consider either implementing the unrolling logic or removing this parameter.

return_sequences=False,
go_backwards=False,
unroll=False,
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot Mar 8, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The unroll parameter is defined but not used in this function. According to the Keras API design guidelines, parameters that are not used should be removed to avoid user confusion and to keep the API surface clean.

time_major=False,
Comment on lines +315 to +316
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot Feb 28, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

These local imports are redundant. keras.src.backend is already imported at the top of the file, and split and concatenate are also imported from keras.src.backend.openvino.numpy at the file's top level. These lines can be removed to improve code clarity.

):
if mask is not None:
raise NotImplementedError("lstm sequence with mask is not supported")
mask = get_ov_output(mask)
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot Mar 8, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

This line is unreachable due to the raise NotImplementedError on the preceding line. It should be removed.

inputs = get_ov_output(inputs)
initial_state_h = get_ov_output(initial_state_h)
initial_state_c = get_ov_output(initial_state_c)
kernel = get_ov_output(kernel)
recurrent_kernel = get_ov_output(recurrent_kernel)
shape = ov_opset.shape_of(inputs, Type.i32).output(0)
if not time_major:
batch_size = ov_opset.gather(
shape,
ov_opset.constant([0], Type.i32),
ov_opset.constant(0, Type.i32),
).output(0)
seq_length = ov_opset.gather(
shape,
ov_opset.constant([1], Type.i32),
ov_opset.constant(0, Type.i32),
).output(0)
else:
batch_size = ov_opset.gather(
shape,
ov_opset.constant([1], Type.i32),
ov_opset.constant(0, Type.i32),
).output(0)
seq_length = ov_opset.gather(
shape,
ov_opset.constant([0], Type.i32),
ov_opset.constant(0, Type.i32),
).output(0)
seq_length_tensor = ov_opset.broadcast(seq_length, batch_size).output(0)
if time_major:
axes = ov_opset.constant([1, 0, 2], Type.i32).output(0)
inputs = ov_opset.transpose(inputs, axes).output(0)
kernel_T = ov_opset.transpose(
kernel, ov_opset.constant([1, 0], Type.i32)
).output(0)
k_i, k_f, k_c, k_o = ov_opset.split(
kernel_T, ov_opset.constant(0, Type.i32), 4
).outputs()
W = ov_opset.concat([k_f, k_i, k_c, k_o], axis=0).output(0)
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot Mar 8, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

critical

The concatenation order for weights W is incorrect. Keras LSTM kernels are ordered for [input, forget, cell, output] gates. The OpenVINO LSTMSequence operation expects the weights for gates in the order [input, output, forget, cell]. The current code concatenates them in the order [forget, input, cell, output]. This should be [input, output, forget, cell].

Suggested change
W = ov_opset.concat([k_f, k_i, k_c, k_o], axis=0).output(0)
W = ov_opset.concat([k_i, k_o, k_f, k_c], axis=0).output(0)

W = ov_opset.unsqueeze(W, ov_opset.constant([0], Type.i32)).output(0)
recurrent_kernel_T = ov_opset.transpose(
recurrent_kernel, ov_opset.constant([1, 0], Type.i32)
).output(0)
r_i, r_f, r_c, r_o = ov_opset.split(
recurrent_kernel_T, ov_opset.constant(0, Type.i32), 4
).outputs()
R = ov_opset.concat([r_f, r_i, r_c, r_o], axis=0).output(0)
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot Mar 8, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

critical

Similar to the kernel weights W, the concatenation order for recurrent weights R is incorrect. It should be [input, output, forget, cell] to match the OpenVINO LSTMSequence op specification.

Suggested change
R = ov_opset.concat([r_f, r_i, r_c, r_o], axis=0).output(0)
R = ov_opset.concat([r_i, r_o, r_f, r_c], axis=0).output(0)

R = ov_opset.unsqueeze(R, ov_opset.constant([0], Type.i32)).output(0)
if bias is not None:
bias = get_ov_output(bias)
b_i, b_f, b_c, b_o = ov_opset.split(
bias, ov_opset.constant(0, Type.i32), 4
).outputs()
B = ov_opset.concat([b_f, b_i, b_c, b_o], axis=0).output(0)
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot Mar 8, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

critical

Similar to the weights, the concatenation order for biases B is incorrect. It should be [input, output, forget, cell] to match the OpenVINO LSTMSequence op specification.

Suggested change
B = ov_opset.concat([b_f, b_i, b_c, b_o], axis=0).output(0)
B = ov_opset.concat([b_i, b_o, b_f, b_c], axis=0).output(0)

else:
W_shape = ov_opset.shape_of(W, Type.i32).output(0)
B_size = ov_opset.gather(
W_shape,
ov_opset.constant([1], Type.i32),
ov_opset.constant(0, Type.i32),
).output(0)
B = ov_opset.broadcast(
ov_opset.constant(0, inputs.get_element_type()), B_size
).output(0)
B = ov_opset.unsqueeze(B, ov_opset.constant([0], Type.i32)).output(0)
initial_state_h = ov_opset.unsqueeze(
initial_state_h, ov_opset.constant([1], Type.i32)
).output(0)
initial_state_c = ov_opset.unsqueeze(
initial_state_c, ov_opset.constant([1], Type.i32)
).output(0)
direction = "forward"
if go_backwards:
direction = "reverse"

def get_activation_name(act):

if act == activations.tanh:
return "tanh"
if act == activations.sigmoid:
return "sigmoid"
if act == activations.relu:
return "relu"
if act == activations.linear:
return "linear"
if hasattr(act, "__name__"):
return act.__name__
return "tanh"
Comment on lines +396 to +408
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot Mar 8, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

This function can be made more concise and maintainable by using a dictionary to map activation functions to their string names. This is a common pattern for this type of logic and would improve readability.

For example:

ACTIVATION_MAP = {
    activations.tanh: "tanh",
    activations.sigmoid: "sigmoid",
    activations.relu: "relu",
    activations.linear: "linear",
}

def get_activation_name(act):
    name = ACTIVATION_MAP.get(act)
    if name:
        return name
    if hasattr(act, "__name__"):
        return act.__name__
    return "tanh"


act_names = [
get_activation_name(recurrent_activation),
get_activation_name(activation),
get_activation_name(activation),
]
hidden_size = (
R.get_partial_shape()[2].get_length()
if R.get_partial_shape().rank.is_static
and R.get_partial_shape()[2].is_static
else -1
)
lstm_node = ov_opset.lstm_sequence(
X=inputs,
initial_hidden_state=initial_state_h,
initial_cell_state=initial_state_c,
sequence_lengths=seq_length_tensor,
W=W,
R=R,
B=B,
hidden_size=hidden_size,
direction=direction,
activations=act_names,
)
Y = lstm_node.output(0)
Ho = lstm_node.output(1)
Co = lstm_node.output(2)
Y = ov_opset.squeeze(Y, ov_opset.constant([1], Type.i32)).output(0)
Ho = ov_opset.squeeze(Ho, ov_opset.constant([1], Type.i32)).output(0)
Co = ov_opset.squeeze(Co, ov_opset.constant([1], Type.i32)).output(0)
if not time_major:
outputs = Y
else:
axes = ov_opset.constant([1, 0, 2], Type.i32).output(0)
outputs = ov_opset.transpose(Y, axes).output(0)
if not return_sequences:
if time_major:
outputs = ov_opset.unsqueeze(
Ho, ov_opset.constant([0], Type.i32)
).output(0)
else:
outputs = ov_opset.unsqueeze(
Ho, ov_opset.constant([1], Type.i32)
).output(0)
Comment on lines +444 to +452
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot Mar 8, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

When return_sequences is False, the caller (keras.layers.RNN) uses the last_output return value and ignores the outputs value. This block of code modifies outputs when return_sequences=False, which makes it dead code. This block should be removed to improve clarity and avoid confusion.

last_output = Ho
return (
OpenVINOKerasTensor(last_output),
OpenVINOKerasTensor(outputs),
[OpenVINOKerasTensor(Ho), OpenVINOKerasTensor(Co)],
)


def gru(*args, **kwargs):
Expand All @@ -305,4 +475,4 @@ def numpy_scan(f, init, xs, reverse=False, mask=None):


def cudnn_ok(*args, **kwargs):
return False
return False
Loading
Loading