Why is the fix TF-only? What makes the problem TF specific?

This is issue with TF backend + Graph Mode + partial_batch_size dataset,i.e when total_samples % batch_size !=0 In which case the dataset adapter returns a strided_slice as batch_shape which is not an integer value. Hence the below check fails with OperatorNotAllowedInGraphError

if mask_shape is not None and mask_shape[:2] != (batch_size, timesteps,): raise ValueError(...)

I am assuming that there is a limitation to set batch_size to a exact integer for this case since last batch will have different batch_size. Hence proposing to skip the evaluation of batch_size for this case i.e "TF backend + Graph Mode + partial_batch_size dataset".

But this change also omits batch_size check for dataset which don't have partial_batch_size also. Hence the below change can also a better fix instead of current commit.

        if backend.backend() == "tensorflow" and not tf.executing_eagerly():
            mask_shape = list(mask_shape)
            mask_shape[0] = tf.get_static_value(mask_shape[0])
            timesteps = tf.get_static_value(timesteps)
            batch_size = tf.get_static_value(batch_size)
            mask_shape = tuple(mask_shape)
        if mask_shape is not None and mask_shape[:2] != (batch_size, timesteps,):
            raise ValueError(
                "`TimeDistributed` Layer should be passed a `mask` of "
                f"shape ({batch_size}, {timesteps}, ...), "
                f"received: mask.shape={mask_shape}"
            )