fix(rnn): bounds check seq_len_max in BlockLSTM ops to prevent OOB access by veerjain-1 · Pull Request #125400 · tensorflow/tensorflow

veerjain-1 · 2026-08-15T19:01:53Z

Resolves #124974. This PR introduces bounds checking for the seq_len_max parameter in the core TensorFlow LSTM and LSTM gradient operations (BlockLSTMOp and BlockLSTMGradOp). Previously, an oversized or negative seq_len_max bypassed validation and triggered a fatal abort (SIGABRT) in Tensor::Slice. The parameter is now validated to ensure it falls within the bounds of [0, timelen].

google-cla · 2026-08-15T19:01:59Z

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

gemini-code-assist

Code Review

This pull request introduces validation checks in BlockLSTMOp and BlockLSTMGradOp within tensorflow/core/kernels/rnn/lstm_ops.cc to ensure that seq_len_max is non-negative and does not exceed timelen. There are no review comments provided, so I have no feedback to evaluate.

…cess

dmiltr3

Please add unit tests covering seq_len_max bounds validation for both BlockLSTM and BlockLSTMGrad in tensorflow/python/kernel_tests/nn_ops/rnn_cell_test.py to ensure invalid bounds (seq_len_max < 0 and seq_len_max > timelen) raise InvalidArgumentError and valid boundary values (0, timelen) execute as expected.

Suggested tests to add to rnn_cell_test.py:

  @test_util.run_in_graph_and_eager_modes
  def testBlockLSTMSeqLenMaxBounds(self):
    timelen, batch_size, input_size, cell_size = 4, 2, 3, 5
    x = constant_op.constant(0.1, shape=[timelen, batch_size, input_size], dtype=dtypes.float32)
    cs_prev = constant_op.constant(0.1, shape=[batch_size, cell_size], dtype=dtypes.float32)
    h_prev = constant_op.constant(0.1, shape=[batch_size, cell_size], dtype=dtypes.float32)
    w = constant_op.constant(0.1, shape=[input_size + cell_size, 4 * cell_size], dtype=dtypes.float32)
    wci = constant_op.constant(0.1, shape=[cell_size], dtype=dtypes.float32)
    wcf = constant_op.constant(0.1, shape=[cell_size], dtype=dtypes.float32)
    wco = constant_op.constant(0.1, shape=[cell_size], dtype=dtypes.float32)
    b = constant_op.constant(0.1, shape=[4 * cell_size], dtype=dtypes.float32)

    # Valid boundary cases: seq_len_max == 0 and seq_len_max == timelen
    for valid_seq_len_max in [0, 2, timelen]:
      res = self.evaluate(
          gen_rnn_ops.BlockLSTM(
              seq_len_max=constant_op.constant(valid_seq_len_max, dtype=dtypes.int64),
              x=x,
              cs_prev=cs_prev,
              h_prev=h_prev,
              w=w,
              wci=wci,
              wcf=wcf,
              wco=wco,
              b=b,
              forget_bias=0.0,
              cell_clip=-1.0,
              use_peephole=False,
          )
      )
      self.assertEqual(len(res), 7)

    # Invalid cases: seq_len_max < 0 or seq_len_max > timelen
    for invalid_seq_len_max in [-1, timelen + 1, 100]:
      with self.assertRaisesRegex(
          (ValueError, errors_impl.InvalidArgumentError),
          r"seq_len_max must be between 0 and",
      ):
        self.evaluate(
            gen_rnn_ops.BlockLSTM(
                seq_len_max=constant_op.constant(invalid_seq_len_max, dtype=dtypes.int64),
                x=x,
                cs_prev=cs_prev,
                h_prev=h_prev,
                w=w,
                wci=wci,
                wcf=wcf,
                wco=wco,
                b=b,
                forget_bias=0.0,
                cell_clip=-1.0,
                use_peephole=False,
            )
        )

  @test_util.run_in_graph_and_eager_modes
  def testBlockLSTMGradSeqLenMaxBounds(self):
    timelen, batch_size, input_size, cell_size = 4, 2, 3, 5
    x = constant_op.constant(0.1, shape=[timelen, batch_size, input_size], dtype=dtypes.float32)
    cs_prev = constant_op.constant(0.1, shape=[batch_size, cell_size], dtype=dtypes.float32)
    h_prev = constant_op.constant(0.1, shape=[batch_size, cell_size], dtype=dtypes.float32)
    w = constant_op.constant(0.1, shape=[input_size + cell_size, 4 * cell_size], dtype=dtypes.float32)
    wci = constant_op.constant(0.1, shape=[cell_size], dtype=dtypes.float32)
    wcf = constant_op.constant(0.1, shape=[cell_size], dtype=dtypes.float32)
    wco = constant_op.constant(0.1, shape=[cell_size], dtype=dtypes.float32)
    b = constant_op.constant(0.1, shape=[4 * cell_size], dtype=dtypes.float32)
    i = constant_op.constant(0.1, shape=[timelen, batch_size, cell_size], dtype=dtypes.float32)
    cs = constant_op.constant(0.1, shape=[timelen, batch_size, cell_size], dtype=dtypes.float32)
    f = constant_op.constant(0.1, shape=[timelen, batch_size, cell_size], dtype=dtypes.float32)
    o = constant_op.constant(0.1, shape=[timelen, batch_size, cell_size], dtype=dtypes.float32)
    ci = constant_op.constant(0.1, shape=[timelen, batch_size, cell_size], dtype=dtypes.float32)
    co = constant_op.constant(0.1, shape=[timelen, batch_size, cell_size], dtype=dtypes.float32)
    h = constant_op.constant(0.1, shape=[timelen, batch_size, cell_size], dtype=dtypes.float32)
    cs_grad = constant_op.constant(0.1, shape=[timelen, batch_size, cell_size], dtype=dtypes.float32)
    h_grad = constant_op.constant(0.1, shape=[timelen, batch_size, cell_size], dtype=dtypes.float32)

    # Valid boundary cases: seq_len_max == 0 and seq_len_max == timelen
    for valid_seq_len_max in [0, 2, timelen]:
      res = self.evaluate(
          gen_rnn_ops.BlockLSTMGrad(
              seq_len_max=constant_op.constant(valid_seq_len_max, dtype=dtypes.int64),
              x=x,
              cs_prev=cs_prev,
              h_prev=h_prev,
              w=w,
              wci=wci,
              wcf=wcf,
              wco=wco,
              b=b,
              i=i,
              cs=cs,
              f=f,
              o=o,
              ci=ci,
              co=co,
              h=h,
              cs_grad=cs_grad,
              h_grad=h_grad,
              use_peephole=False,
          )
      )
      self.assertEqual(len(res), 8)

    # Invalid cases: seq_len_max < 0 or seq_len_max > timelen
    for invalid_seq_len_max in [-1, timelen + 1, 100]:
      with self.assertRaisesRegex(
          (ValueError, errors_impl.InvalidArgumentError),
          r"seq_len_max must be between 0 and",
      ):
        self.evaluate(
            gen_rnn_ops.BlockLSTMGrad(
                seq_len_max=constant_op.constant(invalid_seq_len_max, dtype=dtypes.int64),
                x=x,
                cs_prev=cs_prev,
                h_prev=h_prev,
                w=w,
                wci=wci,
                wcf=wcf,
                wco=wco,
                b=b,
                i=i,
                cs=cs,
                f=f,
                o=o,
                ci=ci,
                co=co,
                h=h,
                cs_grad=cs_grad,
                h_grad=h_grad,
                use_peephole=False,
            )
        )

veerjain-1 · 2026-08-20T17:27:26Z

Hi team! Just a gentle ping on this PR. Let me know if there's anything else you need from my side or if any changes are required. Thanks!

dmiltr3

Please add unit tests covering seq_len_max bounds validation for both BlockLSTM and BlockLSTMGrad in tensorflow/python/kernel_tests/nn_ops/rnn_cell_test.py to ensure invalid bounds (seq_len_max < 0 and seq_len_max > timelen) raise InvalidArgumentError and valid boundary values (0, timelen) execute as expected.

Suggested tests to add to rnn_cell_test.py:

  @test_util.run_in_graph_and_eager_modes
  def testBlockLSTMSeqLenMaxBounds(self):
    timelen, batch_size, input_size, cell_size = 4, 2, 3, 5
    x = constant_op.constant(0.1, shape=[timelen, batch_size, input_size], dtype=dtypes.float32)
    cs_prev = constant_op.constant(0.1, shape=[batch_size, cell_size], dtype=dtypes.float32)
    h_prev = constant_op.constant(0.1, shape=[batch_size, cell_size], dtype=dtypes.float32)
    w = constant_op.constant(0.1, shape=[input_size + cell_size, 4 * cell_size], dtype=dtypes.float32)
    wci = constant_op.constant(0.1, shape=[cell_size], dtype=dtypes.float32)
    wcf = constant_op.constant(0.1, shape=[cell_size], dtype=dtypes.float32)
    wco = constant_op.constant(0.1, shape=[cell_size], dtype=dtypes.float32)
    b = constant_op.constant(0.1, shape=[4 * cell_size], dtype=dtypes.float32)

    # Valid boundary cases: seq_len_max == 0 and seq_len_max == timelen
    for valid_seq_len_max in [0, 2, timelen]:
      res = self.evaluate(
          gen_rnn_ops.BlockLSTM(
              seq_len_max=constant_op.constant(valid_seq_len_max, dtype=dtypes.int64),
              x=x,
              cs_prev=cs_prev,
              h_prev=h_prev,
              w=w,
              wci=wci,
              wcf=wcf,
              wco=wco,
              b=b,
              forget_bias=0.0,
              cell_clip=-1.0,
              use_peephole=False,
          )
      )
      self.assertEqual(len(res), 7)

    # Invalid cases: seq_len_max < 0 or seq_len_max > timelen
    for invalid_seq_len_max in [-1, timelen + 1, 100]:
      with self.assertRaisesRegex(
          (ValueError, errors_impl.InvalidArgumentError),
          r"seq_len_max must be between 0 and",
      ):
        self.evaluate(
            gen_rnn_ops.BlockLSTM(
                seq_len_max=constant_op.constant(invalid_seq_len_max, dtype=dtypes.int64),
                x=x,
                cs_prev=cs_prev,
                h_prev=h_prev,
                w=w,
                wci=wci,
                wcf=wcf,
                wco=wco,
                b=b,
                forget_bias=0.0,
                cell_clip=-1.0,
                use_peephole=False,
            )
        )

  @test_util.run_in_graph_and_eager_modes
  def testBlockLSTMGradSeqLenMaxBounds(self):
    timelen, batch_size, input_size, cell_size = 4, 2, 3, 5
    x = constant_op.constant(0.1, shape=[timelen, batch_size, input_size], dtype=dtypes.float32)
    cs_prev = constant_op.constant(0.1, shape=[batch_size, cell_size], dtype=dtypes.float32)
    h_prev = constant_op.constant(0.1, shape=[batch_size, cell_size], dtype=dtypes.float32)
    w = constant_op.constant(0.1, shape=[input_size + cell_size, 4 * cell_size], dtype=dtypes.float32)
    wci = constant_op.constant(0.1, shape=[cell_size], dtype=dtypes.float32)
    wcf = constant_op.constant(0.1, shape=[cell_size], dtype=dtypes.float32)
    wco = constant_op.constant(0.1, shape=[cell_size], dtype=dtypes.float32)
    b = constant_op.constant(0.1, shape=[4 * cell_size], dtype=dtypes.float32)
    i = constant_op.constant(0.1, shape=[timelen, batch_size, cell_size], dtype=dtypes.float32)
    cs = constant_op.constant(0.1, shape=[timelen, batch_size, cell_size], dtype=dtypes.float32)
    f = constant_op.constant(0.1, shape=[timelen, batch_size, cell_size], dtype=dtypes.float32)
    o = constant_op.constant(0.1, shape=[timelen, batch_size, cell_size], dtype=dtypes.float32)
    ci = constant_op.constant(0.1, shape=[timelen, batch_size, cell_size], dtype=dtypes.float32)
    co = constant_op.constant(0.1, shape=[timelen, batch_size, cell_size], dtype=dtypes.float32)
    h = constant_op.constant(0.1, shape=[timelen, batch_size, cell_size], dtype=dtypes.float32)
    cs_grad = constant_op.constant(0.1, shape=[timelen, batch_size, cell_size], dtype=dtypes.float32)
    h_grad = constant_op.constant(0.1, shape=[timelen, batch_size, cell_size], dtype=dtypes.float32)

    # Valid boundary cases: seq_len_max == 0 and seq_len_max == timelen
    for valid_seq_len_max in [0, 2, timelen]:
      res = self.evaluate(
          gen_rnn_ops.BlockLSTMGrad(
              seq_len_max=constant_op.constant(valid_seq_len_max, dtype=dtypes.int64),
              x=x,
              cs_prev=cs_prev,
              h_prev=h_prev,
              w=w,
              wci=wci,
              wcf=wcf,
              wco=wco,
              b=b,
              i=i,
              cs=cs,
              f=f,
              o=o,
              ci=ci,
              co=co,
              h=h,
              cs_grad=cs_grad,
              h_grad=h_grad,
              use_peephole=False,
          )
      )
      self.assertEqual(len(res), 8)

    # Invalid cases: seq_len_max < 0 or seq_len_max > timelen
    for invalid_seq_len_max in [-1, timelen + 1, 100]:
      with self.assertRaisesRegex(
          (ValueError, errors_impl.InvalidArgumentError),
          r"seq_len_max must be between 0 and",
      ):
        self.evaluate(
            gen_rnn_ops.BlockLSTMGrad(
                seq_len_max=constant_op.constant(invalid_seq_len_max, dtype=dtypes.int64),
                x=x,
                cs_prev=cs_prev,
                h_prev=h_prev,
                w=w,
                wci=wci,
                wcf=wcf,
                wco=wco,
                b=b,
                i=i,
                cs=cs,
                f=f,
                o=o,
                ci=ci,
                co=co,
                h=h,
                cs_grad=cs_grad,
                h_grad=h_grad,
                use_peephole=False,
            )
        )

veerjain-1 · 2026-08-23T19:24:08Z

Hi @dmiltr3, I have added the requested unit tests covering seq_len_max bounds validation in rnn_cell_test.py and resolved the merge conflicts. Could you please take another look? Thanks!

dmiltr3

CI Failure Investigation & Required Test Fix

The GitHub CI checks failed across multiple platforms:

Root Cause

In tensorflow/core/kernels/rnn/lstm_ops.cc, consolidating the bounds check error message into "seq_len_max must be between 0 and ..." caused 4 existing unit tests in tensorflow/python/ops/rnn_grad_test.py to fail (testBlockLSTMSeqLenMaxNegative, testBlockLSTMSeqLenMaxTooLarge, testBlockLSTMGradSeqLenMaxNegative, testBlockLSTMGradSeqLenMaxTooLarge) because their assertion regex expects "seq_len_max must be <= timelen" or "seq_len_max must be >= 0".

Proposed Fix

Please update the expected error message regex in tensorflow/python/ops/rnn_grad_test.py:

diff --git a/tensorflow/python/ops/rnn_grad_test.py b/tensorflow/python/ops/rnn_grad_test.py
index a5d3e09..b61c94d 100644
--- a/tensorflow/python/ops/rnn_grad_test.py
+++ b/tensorflow/python/ops/rnn_grad_test.py
@@ -120,7 +120,7 @@ class RNNGradTest(test.TestCase):
   def testBlockLSTMSeqLenMaxTooLarge(self):
     w, b, x, cs_prev, h_prev, w_peephole = self._block_lstm_inputs()
     with self.assertRaisesRegex(
-        errors_impl.InvalidArgumentError, "seq_len_max must be <= timelen"
+        errors_impl.InvalidArgumentError, r"seq_len_max must be between 0 and"
     ):
       self.evaluate(
           self._block_lstm(w, b, x, cs_prev, h_prev, w_peephole, seq_len_max=10)
@@ -128,7 +128,7 @@ class RNNGradTest(test.TestCase):
   def testBlockLSTMSeqLenMaxNegative(self):
     w, b, x, cs_prev, h_prev, w_peephole = self._block_lstm_inputs()
     with self.assertRaisesRegex(
-        errors_impl.InvalidArgumentError, "seq_len_max must be >= 0"
+        errors_impl.InvalidArgumentError, r"seq_len_max must be between 0 and"
     ):
       self.evaluate(
           self._block_lstm(w, b, x, cs_prev, h_prev, w_peephole, seq_len_max=-1)
@@ -138,7 +138,7 @@ class RNNGradTest(test.TestCase):
   def testBlockLSTMGradSeqLenMaxTooLarge(self):
     with self.assertRaisesRegex(
-        errors_impl.InvalidArgumentError, "seq_len_max must be <= timelen"
+        errors_impl.InvalidArgumentError, r"seq_len_max must be between 0 and"
     ):
       self.evaluate(self._block_lstm_grad(seq_len_max=10))
 
@@ -142,7 +142,7 @@ class RNNGradTest(test.TestCase):
   def testBlockLSTMGradSeqLenMaxNegative(self):
     with self.assertRaisesRegex(
-        errors_impl.InvalidArgumentError, "seq_len_max must be >= 0"
+        errors_impl.InvalidArgumentError, r"seq_len_max must be between 0 and"
     ):
       self.evaluate(self._block_lstm_grad(seq_len_max=-1))

veerjain-1 · 2026-08-24T18:29:09Z

Hi @dmiltr3, thanks for investigating the CI failure! I have updated the expected error message regex in rnn_grad_test.py as suggested, and also fixed some PyLint line-length errors in the unit tests I added. Let me know if you need anything else!

dmiltr3

All tests and verification checks have passed. Changes look good.

veerjain-1 · 2026-08-25T14:43:03Z

I've applied the requested test fix for the regex. Could you please take another look and resolve the block if everything looks good? Also, I noticed some CI failures (build-windows-x86 and build-linux-x86-cpu) - do those look like unrelated flaky tests, or is there something else I need to address on my end?

nithyak0204 · 2026-08-26T08:43:54Z

Hi @veerjain-1, this PR is waiting for internal approvals. We are working on it. Thank you!

google-ml-butler Bot added the size:S CL Change Size: Small label Aug 15, 2026

google-ml-butler Bot assigned gbaned Aug 15, 2026

gemini-code-assist Bot reviewed Aug 15, 2026

View reviewed changes

veerjain-1 force-pushed the fix-blocklstmgrad-oob branch from b3fd4a8 to dbb38dc Compare August 15, 2026 19:11

fix(rnn): bounds check seq_len_max in BlockLSTM ops to prevent OOB ac…

86d7965

…cess

veerjain-1 force-pushed the fix-blocklstmgrad-oob branch from dbb38dc to 86d7965 Compare August 15, 2026 19:15

nithyak0204 requested a review from a team August 17, 2026 07:42

google-ml-butler Bot added the awaiting review Pull request awaiting review label Aug 17, 2026

nithyak0204 added comp:core issues related to core part of tensorflow prtype:bugfix PR to fix a bug labels Aug 17, 2026

nithyak0204 added this to PR Queue Aug 17, 2026

github-project-automation Bot moved this to Assigned Reviewer in PR Queue Aug 17, 2026

dmiltr3 self-requested a review August 17, 2026 14:55

dmiltr3 requested changes Aug 17, 2026

View reviewed changes

github-project-automation Bot moved this from Assigned Reviewer to Reviewer Requested Changes in PR Queue Aug 17, 2026

nithyak0204 added stat:awaiting response Status - Awaiting response from author and removed awaiting review Pull request awaiting review labels Aug 18, 2026

google-ml-butler Bot removed the stat:awaiting response Status - Awaiting response from author label Aug 20, 2026

dmiltr3 requested changes Aug 20, 2026

View reviewed changes

veerjain-1 added 2 commits August 23, 2026 11:32

test: add bounds validation tests for BlockLSTM seq_len_max

c841dac

Merge master and resolve conflicts

592bd0d

fix(test): wrap long lines in rnn_cell_test.py to pass PyLint

f8fbb96

dmiltr3 requested changes Aug 23, 2026

View reviewed changes

fix(test): update expected error message for seq_len_max bounds check

40ad276

dmiltr3 approved these changes Aug 24, 2026

View reviewed changes

google-ml-butler Bot added the kokoro:force-run Tests on submitted change label Aug 24, 2026

google-ml-butler Bot added the ready to pull PR ready for merge process label Aug 24, 2026

github-project-automation Bot moved this from Reviewer Requested Changes to Approved by Reviewer in PR Queue Aug 24, 2026

vishwakt mentioned this pull request Aug 24, 2026

Validate input shapes in the LSTM gradient kernels #126028

Merged

copybara-service Bot merged commit 786097d into tensorflow:master Sep 1, 2026
15 of 17 checks passed

github-project-automation Bot moved this from Approved by Reviewer to Merged in PR Queue Sep 1, 2026

Ashutosh0x mentioned this pull request Sep 9, 2026

fix: validate input rank before reading dims in LSTMBlockCell kernels #124034

Merged

Navigation Menu

Repository navigation

fix(rnn): bounds check seq_len_max in BlockLSTM ops to prevent OOB access - #125400

fix(rnn): bounds check seq_len_max in BlockLSTM ops to prevent OOB access#125400
copybara-service[bot] merged 5 commits into
tensorflow:masterfrom
veerjain-1:fix-blocklstmgrad-oob

veerjain-1 commented Aug 15, 2026

Uh oh!

google-cla Bot commented Aug 15, 2026

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

dmiltr3 left a comment

Uh oh!

veerjain-1 commented Aug 20, 2026

Uh oh!

dmiltr3 left a comment

Uh oh!

veerjain-1 commented Aug 23, 2026

Uh oh!

dmiltr3 left a comment

Uh oh!

veerjain-1 commented Aug 24, 2026

Uh oh!

dmiltr3 left a comment

Uh oh!

veerjain-1 commented Aug 25, 2026

Uh oh!

nithyak0204 commented Aug 26, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Conversation

veerjain-1 commented Aug 15, 2026

Uh oh!

google-cla Bot commented Aug 15, 2026

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

dmiltr3 left a comment

Choose a reason for hiding this comment

Uh oh!

veerjain-1 commented Aug 20, 2026

Uh oh!

dmiltr3 left a comment

Choose a reason for hiding this comment

Uh oh!

veerjain-1 commented Aug 23, 2026

Uh oh!

dmiltr3 left a comment

Choose a reason for hiding this comment

CI Failure Investigation & Required Test Fix

Root Cause

Proposed Fix

Uh oh!

veerjain-1 commented Aug 24, 2026

Uh oh!

dmiltr3 left a comment

Choose a reason for hiding this comment

Uh oh!

veerjain-1 commented Aug 25, 2026

Uh oh!

nithyak0204 commented Aug 26, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants