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 9 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
194 changes: 183 additions & 11 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,8 @@
from openvino import Model
from openvino import Type

from keras.src import activations
from keras.src import ops
Copy link
Collaborator

@hertschuh hertschuh Mar 4, 2026

Choose a reason for hiding this comment

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

You cannot import ops like this here, this is what's creating the circular imports.

But they're not even used.

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 +107,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 +132,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 +146,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 +174,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 +198,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 +234,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 +301,163 @@ 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):
from keras.src import activations
Copy link
Collaborator

@hertschuh hertschuh Mar 4, 2026

Choose a reason for hiding this comment

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

activations are already imported at the top, you don't need this import.


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 @@ -304,5 +476,5 @@ def numpy_scan(f, init, xs, reverse=False, mask=None):
)


def cudnn_ok(*args, **kwargs):
return False
def cudnn_ok(*args , **kwargs):
return False
Copy link
Collaborator

@hertschuh hertschuh Mar 4, 2026

Choose a reason for hiding this comment

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

Revert this change file, it's failing the code format check.

Copy link
Contributor Author

@andersendsa andersendsa Mar 5, 2026

Choose a reason for hiding this comment

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

Hey @hertschuh just wanted to double-check—do you want me to revert the whole file back to the previous commit?

Copy link
Collaborator

@hertschuh hertschuh Mar 5, 2026

Choose a reason for hiding this comment

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

Oh, no, just the last 2 lines of this file.

2 changes: 1 addition & 1 deletion keras/src/initializers/random_initializers.py
View file
Open in desktop
Copy link
Contributor

@goyaladitya05 goyaladitya05 Mar 1, 2026

Choose a reason for hiding this comment

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

I'm confused by the inclusion of these changes in this PR. keras/src/initializers/ is a global utility and should remain backend-agnostic. Unless there is a specific reason these initializers need to be changed for all of Keras, I'd suggest reverting the changes to this file.

Copy link
Collaborator

@hertschuh hertschuh Mar 4, 2026

Choose a reason for hiding this comment

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

+1

Why was this change needed?

Copy link
Collaborator

@hertschuh hertschuh Mar 4, 2026

Choose a reason for hiding this comment

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

Revert this whole file, it's failing the code format check.

Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@

from keras.src import ops
from keras.src.api_export import keras_export
from keras.src.backend import random
from keras.src.backend import random
from keras.src.initializers.initializer import Initializer
from keras.src.saving import serialization_lib

Expand Down
2 changes: 1 addition & 1 deletion keras/src/ops/__init__.py
View file
Open in desktop
Copy link
Contributor

@goyaladitya05 goyaladitya05 Mar 1, 2026

Choose a reason for hiding this comment

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

Manually overriding dir and getattr to force-import the entire library seems unnecessary. We only make changes to the files involving the backend, for a backend specific implementation. Is there a reason for this?

Copy link
Collaborator

@hertschuh hertschuh Mar 4, 2026

Choose a reason for hiding this comment

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

+1

Revert this.

Copy link
Collaborator

@hertschuh hertschuh Mar 4, 2026

Choose a reason for hiding this comment

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

Revert this whole file, it's failing the code format check.

Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
# from keras.src.ops.numpy import Matmul, matmul
# from keras.src.ops.numpy import Add, add
# from keras.src.ops.numpy import Multiply, multiply

from keras.src.backend import cast
from keras.src.backend import cond
from keras.src.backend import is_tensor
Expand All @@ -14,3 +13,4 @@
from keras.src.ops.math import * # noqa: F403
from keras.src.ops.nn import * # noqa: F403
from keras.src.ops.numpy import * # noqa: F403

Loading