Fix for #18941 (#18968)

This allows the static shape to propagate in RNNs in the case when the number of time steps is fixed and known at build time.
keras-team · Dec 19, 2023 · 4a4a139 · 4a4a139
1 parent 2fbc05d
commit 4a4a139
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/keras/backend/tensorflow/rnn.py b/keras/backend/tensorflow/rnn.py
@@ -90,7 +90,9 @@ def swap_batch_timestep(input_t):
 
     flattened_inputs = tree.flatten(inputs)
     time_steps = flattened_inputs[0].shape[0]
-    time_steps_t = tf.shape(flattened_inputs[0])[0]
+    time_steps_t = (
+        tf.shape(flattened_inputs[0])[0] if time_steps is None else time_steps
+    )
 
     for input_ in flattened_inputs:
         input_.shape.with_rank_at_least(3)