Skip to content

[OpenVINO] Implements LSTM and GRU, fixes RNN shape bug#22436

Open
goyaladitya05 wants to merge 7 commits intokeras-team:masterfrom
goyaladitya05:openvino/rnn-lstm-gru
Open

[OpenVINO] Implements LSTM and GRU, fixes RNN shape bug#22436
goyaladitya05 wants to merge 7 commits intokeras-team:masterfrom
goyaladitya05:openvino/rnn-lstm-gru

Conversation

@goyaladitya05
Copy link
Contributor

@goyaladitya05 goyaladitya05 commented Mar 15, 2026

This PR implements lstm() and gru() for the OpenVINO backend using ov_opset.lstm_sequence and gru_sequence, replacing the backend.rnn() fallback for the common case.

This PR also fixes a pre-existing shape bug in backend.rnn() that caused all RNN-based layers (LSTM, GRU, SimpleRNN, StackedRNN) to fail or produce wrong results with implementation=2 (the default).

Fix: Declared body params as [1, batch, features], squeeze axis 0 before calling step_function, and unsqueeze the step output back to [1, batch, units] before get_concatenated_slices.

LSTM and GRU implementation

Uses ov_opset.lstm_sequence, gru_sequence when all of the following hold:

  • activations are tanh + sigmoid
  • unroll=False
  • bias is provided
  • mask is None
  • GRU only: reset_after=True

Weight layout differences handled:

  • LSTM: Keras gate order [i, f, c, o] vs OV [f, i, c, o] (reordered via split/concat)
  • GRU: same gate order [z, r, h] in both; bias [2, 3*units] vs OV [1, 4*units] = [b_z+rb_z, b_r+rb_r, b_h, rb_h] for linear_before_reset=True

Removed numpy_scan from rnn.py. It was implemented to throw NotImplementedError, and was not present at all for other backends (Tensorflow, Torch, JAX).

Removed unstack from rnn.py. Tensorflow and Torch have it being implemented in core.py, and not in rnn.py.

Tests

Enabled 30 tests releted to RNN,LSTM and GRU from excluded_concrete_tests.txt

Closes: #22435, Closes: openvinotoolkit/openvino/issues/34403, Closes: openvinotoolkit/openvino/issues/34408

Note

After this PR, rnn.py is at parity with other backends. RNN, LSTM, GRU, and cudnn_ok are all implemented.

@gemini-code-assist
Copy link
Contributor

Summary of Changes

Hello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly enhances the OpenVINO backend by providing native implementations for LSTM and GRU recurrent neural network layers, moving away from a generic fallback. Concurrently, it resolves a critical shape-related bug within the existing RNN backend, ensuring the stability and correctness of all RNN-based operations. These changes improve performance and reliability for recurrent models on the OpenVINO platform.

Highlights

  • OpenVINO LSTM and GRU Implementation: Implemented native lstm() and gru() functions for the OpenVINO backend, leveraging ov_opset.lstm_sequence and gru_sequence to replace the generic backend.rnn() fallback for common use cases.
  • RNN Shape Bug Fix: Addressed a pre-existing shape bug in backend.rnn() that caused failures or incorrect results for all RNN-based layers (LSTM, GRU, SimpleRNN, StackedRNN) when implementation=2 was used. The fix involved declaring body parameters with a leading [1], squeezing axis 0 before calling the step function, and unsqueezing the step output back to [1, batch, units].
  • Conditional OpenVINO Backend Usage: The new ov_opset.lstm_sequence and gru_sequence operations are utilized when specific conditions are met, including tanh and sigmoid activations, unroll=False, bias presence, mask=None, and reset_after=True for GRU.
  • Weight Layout Handling: Handled differences in weight layouts between Keras and OpenVINO for LSTM (Keras [i, f, c, o] vs OV [f, i, c, o]) and GRU biases (Keras [2, 3*units] vs OV [1, 4*units]).
  • Test Suite Update: Removed 12 tests related to LSTM and GRU from excluded_concrete_tests.txt, indicating that these tests are now expected to pass with the new OpenVINO implementations.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog
  • keras/src/backend/openvino/excluded_concrete_tests.txt
    • Removed 12 tests for GRU and LSTM, as they are now expected to pass with the new OpenVINO implementations.
  • keras/src/backend/openvino/rnn.py
    • Fixed a shape bug in the _rnn_loop_body function by adjusting parameter shapes and adding squeeze/unsqueeze operations.
    • Implemented the lstm function using ov_opset.lstm_sequence, including logic for gate reordering and bias transformation.
    • Implemented the gru function using ov_opset.gru_sequence, including logic for bias transformation to match OpenVINO's expected format.
    • Added helper functions _reorder_gates and _seq_lengths to facilitate the new LSTM and GRU implementations.
Activity
Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature Command Description
Code Review /gemini review Performs a code review for the current pull request in its current state.
Pull Request Summary /gemini summary Provides a summary of the current pull request in its current state.
Comment @gemini-code-assist Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help /gemini help Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

  1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution.

Copy link
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

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

Code Review

This pull request introduces OpenVINO backend implementations for lstm and gru, which is a significant feature addition. It also includes a crucial fix for a shape-related bug in the generic rnn implementation. My review has identified a critical issue in the new lstm implementation concerning bias handling that requires attention. Additionally, I've provided a few suggestions to enhance code quality and enable masking support, which appears to be mostly implemented but is currently disabled.

@codecov-commenter
Copy link

codecov-commenter commented Mar 15, 2026

Codecov Report

❌ Patch coverage is 98.90110% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 83.24%. Comparing base (dc3bdb1) to head (d208cfb).

Files with missing lines Patch % Lines
keras/src/backend/openvino/rnn.py 98.90% 1 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##           master   #22436      +/-   ##
==========================================
- Coverage   84.93%   83.24%   -1.70%     
==========================================
  Files         596      596              
  Lines       66936    67019      +83     
  Branches    10449    10451       +2     
==========================================
- Hits        56855    55790    -1065     
- Misses       7292     8549    +1257     
+ Partials     2789     2680     -109     
Flag Coverage Δ
keras 83.07% <98.90%> (-1.70%) ⬇️
keras-jax 60.29% <4.39%> (-2.04%) ⬇️
keras-numpy 54.55% <0.00%> (-0.17%) ⬇️
keras-openvino 50.80% <98.90%> (+0.62%) ⬆️
keras-tensorflow 61.54% <4.39%> (-2.04%) ⬇️
keras-torch 60.36% <4.39%> (-2.03%) ⬇️
keras.applications ?
keras.applications-jax ?
keras.applications-numpy ?
keras.applications-openvino ?
keras.applications-tensorflow ?
keras.applications-torch ?

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@goyaladitya05 goyaladitya05 marked this pull request as ready for review March 15, 2026 04:44
SpectralNormalizationTest::test_apply_layer
StackedRNNTest::test_correctness_single_state_stack
StackedRNNTest::test_correctness_two_states_stack
StackedRNNTest::test_return_state_stacked_lstm_cell
Copy link
Contributor Author

@goyaladitya05 goyaladitya05 Mar 15, 2026

Choose a reason for hiding this comment

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

StackedRNNTest::test_return_state_stacked_lstm_cell can be enabled once qr op is implemented in linalg.py. It won't pass right now.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

3 participants