Raise exception on batch_size mismatch for stateful RNNs (#21742)

* fix stateful rnn (pulling from existing PR that closed)

* Update keras/src/layers/rnn/rnn.py

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

* ensure we don't recompute expected batch size for each call

* address comments from francois

---------

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Author: danielenricocahall

keras/src/layers/rnn/rnn.py

@@ -212,6 +212,7 @@ def __init__(
         self.supports_masking = True
         self.input_spec = None
         self.states = None
+        self._expected_batch_size = None
 
         state_size = getattr(self.cell, "state_size", None)
         if state_size is None:
@@ -283,6 +284,9 @@ def build(self, sequences_shape, initial_state_shape=None):
                         f"batch size: sequence.shape={sequences_shape}"
                     )
                 self._create_state_variables(sequences_shape[0])
+                self._expected_batch_size = ops.shape(
+                    tree.flatten(self.states)[0]
+                )[0]
 
     @tracking.no_automatic_dependency_tracking
     def _create_state_variables(self, batch_size):
@@ -382,6 +386,21 @@ def call(
                 initial_state = self.get_initial_state(
                     batch_size=ops.shape(sequences)[0]
                 )
+        if self.stateful:
+            actual_batch_size = ops.shape(sequences)[0]
+            if (
+                self._expected_batch_size is not None
+                and actual_batch_size is not None
+                and actual_batch_size != self._expected_batch_size
+            ):
+                raise ValueError(
+                    f"If an RNN is stateful, the batch size of the "
+                    f"input sequences must be the same as the batch "
+                    f"size of the initial state. \n"
+                    f"- Expected batch size: {self._expected_batch_size}\n"
+                    f"- Received batch size: {actual_batch_size}"
+                )
+
         # RNN expect the states in a list, even if single state.
         if not tree.is_nested(initial_state):
             initial_state = [initial_state]

keras/src/layers/rnn/rnn_test.py

@@ -381,4 +381,26 @@ def test_serialization(self):
         layer = layers.RNN(OneStateRNNCell(2), return_sequences=False)
         self.run_class_serialization_test(layer)
 
+    def test_stateful_batch_size_mismatch_raises(self):
+        from keras.src.models import Functional
+
+        batch_size = 4
+        timesteps = 5
+        features = 3
+
+        layer = layers.RNN(TwoStatesRNNCell(2), stateful=True)
+        inputs = layers.Input(
+            shape=(timesteps, features), batch_size=batch_size
+        )
+        model = Functional(inputs, layer(inputs))
+
+        # Call once with correct batch size
+        x = ops.random.uniform(shape=(batch_size, timesteps, features))
+        _ = model(x)
+
+        # Expect ValueError when called with incorrect batch size
+        with self.assertRaisesRegex(ValueError, "batch size"):
+            x_bad = ops.random.uniform(shape=(1, timesteps, features))
+            model(x_bad)
+
     # TODO: test masking