fix time_distributed layer with mask and partial_batch_size (#20765)

* fix time_distributed layer with mask and partial_batch_size

* Fix test fails for non TF backends

* Fix formatting issue

* test case and inline import of TF

* Disable testcase for Numpy backend

* Fix lint error

keras/src/layers/rnn/time_distributed.py

@@ -77,10 +77,29 @@ def call(self, inputs, training=None, mask=None):
         batch_size = input_shape[0]
         timesteps = input_shape[1]
 
-        if mask_shape is not None and mask_shape[:2] != (batch_size, timesteps):
+        # For TF backend with graph mode and `partial_batch_size`, skip
+        # evaluation of `batch_size` as it can be a `strided_slice` and
+        # not a constant.
+        if backend.backend() == "tensorflow":
+            from keras.src.utils.module_utils import tensorflow as tf
+
+            if (
+                not tf.executing_eagerly
+                and mask_shape is not None
+                and mask_shape[1:2] != (timesteps,)
+            ):
+                raise ValueError(
+                    "`TimeDistributed` Layer should be passed a `mask` of "
+                    f"shape ({batch_size}, {timesteps}, ...), "
+                    f"received: mask.shape={mask_shape}"
+                )
+        elif mask_shape is not None and mask_shape[:2] != (
+            batch_size,
+            timesteps,
+        ):
             raise ValueError(
-                "`TimeDistributed` Layer should be passed a `mask` of shape "
-                f"({batch_size}, {timesteps}, ...), "
+                "`TimeDistributed` Layer should be passed a `mask` of "
+                f"shape ({batch_size}, {timesteps}, ...), "
                 f"received: mask.shape={mask_shape}"
             )
 

keras/src/layers/rnn/time_distributed_test.py

@@ -6,6 +6,7 @@
 from keras.src import layers
 from keras.src import ops
 from keras.src import testing
+from keras.src.models import Sequential
 
 
 class TimeDistributedTest(testing.TestCase):
@@ -77,3 +78,24 @@ def call(self, inputs, training=False, mask=None):
             np.array([[[0], [0.22]], [[0.38], [0]], [[0.7], [0.86]]]),
             output,
         )
+
+    @pytest.mark.requires_trainable_backend
+    def test_with_mask_zero(self):
+        model = Sequential(
+            [
+                layers.Input(shape=(20,)),
+                layers.Embedding(input_dim=10, output_dim=5, mask_zero=True),
+                layers.TimeDistributed(
+                    layers.Dense(units=5, activation="softmax")
+                ),
+            ]
+        )
+        model.compile(
+            optimizer="adam",
+            loss="sparse_categorical_crossentropy",
+            metrics=["accuracy"],
+        )
+        X_train = np.random.uniform(1, 10, size=(22, 20))
+        Y_train = np.random.randint(1, 2, size=(22, 20))
+
+        model.fit(X_train, Y_train, epochs=1, batch_size=16)