Repository: keras-team/keras

Issue:
Issue with TimeDistributed + LSTM layer
From [here](https://stackoverflow.com/q/77654253/9215780).

```python
from numpy import array
import keras
from keras.models import Sequential
from keras.layers import Dense
from keras.layers import TimeDistributed
from keras.layers import LSTM

# prepare sequence
length = 5
seq = array([i/float(length) for i in range(length)])
X = seq.reshape(1, length, 1)
y = seq.reshape(1, length, 1)

# define LSTM configuration
n_neurons = length
n_batch = 1
n_epoch = 10

# create LSTM
model = Sequential()
model.add(keras.layers.InputLayer((length, 1)))
model.add(LSTM(n_neurons, return_sequences=True))
model.add(TimeDistributed(Dense(1)))
model.compile(loss='mean_squared_error', optimizer='adam')
model.fit(X, y, epochs=n_epoch, batch_size=n_batch, verbose=2)
```
```
Epoch 1/10
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[6], line 1
----> 1 model.fit(X, y, epochs=n_epoch, batch_size=n_batch, verbose=2)

File /opt/conda/lib/python3.10/site-packages/keras/src/utils/traceback_utils.py:123, in filter_traceback.<locals>.error_handler(*args, **kwargs)
    120     filtered_tb = _process_traceback_frames(e.__traceback__)
    121     # To get the full stack trace, call:
    122     # `keras.config.disable_traceback_filtering()`
--> 123     raise e.with_traceback(filtered_tb) from None
    124 finally:
    125     del filtered_tb

File /opt/conda/lib/python3.10/site-packages/keras/src/backend/common/variables.py:394, in standardize_dtype(dtype)
    391     dtype = str(dtype).split(".")[-1]
    393 if dtype not in ALLOWED_DTYPES:
--> 394     raise ValueError(f"Invalid dtype: {dtype}")
    395 return dtype

ValueError: Exception encountered when calling TimeDistributed.call().

Invalid dtype: <class 'NoneType'>

Arguments received by TimeDistributed.call():
  • inputs=tf.Tensor(shape=(1, None, 5), dtype=float32)
  • training=True
  • mask=None
```

However, compiling the model with eager mode runs properly.



Discussion:
Hi @innat ,

I have replicated the reported error and attached [gist](https://colab.sandbox.google.com/gist/SuryanarayanaY/f07ad210f3e95153156ba8ddcaf3fcd5/18941.ipynb) here.

The reported error occurs when dtype not in `allowed_types` in `keras3`. The inputs `X` and `y` dtypes can be  is `"float64"` after applying `standardize_dtype()` which is OK. It seems the issue with `TimeDistributed` layer output. Needs investigation.

Thanks!

Resolving patch:
diff --git a/keras/backend/tensorflow/rnn.py b/keras/backend/tensorflow/rnn.py
--- 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)


Test patch: