TensorFlow getting all states from a `RNN`
ML System Design practice on Codemia
Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.
Introduction
When people ask for "all states" from an RNN in TensorFlow, they usually mean one of two different things. They either want the output at every time step, or they want the final hidden state values returned separately for reuse in another call.
return_sequences Gives You the Per-Step Outputs
For most Keras RNN layers, the easiest way to get all time-step outputs is return_sequences=True. That tells the layer not to collapse the sequence down to only the last output.
The shape means:
- batch size:
2 - time steps:
5 - hidden size:
4
In many discussions, these per-step outputs are what people actually mean by "all states."
return_state Gives You the Final State
If you also need the final state explicitly, add return_state=True:
For SimpleRNN, the final state usually matches the last output step. For more complex cells such as LSTM, the returned state structure is richer.
LSTM and GRU Return Different State Shapes
GRU returns one final state tensor. LSTM returns two final states: hidden state and cell state.
This is important because asking for "all states" in an LSTM can mean:
- all output vectors across the sequence
- the final hidden state
- the final cell state
Those are related, but not identical.
If You Need Every Time Step for a Later Layer
Set return_sequences=True when another sequence-processing layer follows. Without it, the next RNN receives only one vector instead of the full sequence.
The first recurrent layer keeps the full time axis alive so the second recurrent layer can read all steps.
If You Need Intermediate State Tensors for Analysis
For many analysis tasks, the sequence output is enough. You can inspect hidden representations over time directly from the returned tensor:
That gives you one vector per time step. For visualization, debugging, or attention-style downstream logic, this is usually the right representation to keep.
Common Confusion: Output vs State
Keras uses both terms, but they are not always interchangeable.
For SimpleRNN:
- last output and final state are effectively the same value
For LSTM:
- last output corresponds to the final hidden state
- final cell state is separate and is not equal to the output sequence tensor
This is why code that works for SimpleRNN can become confusing when switched to LSTM.
Stateful RNNs Are a Different Feature
Do not confuse return_state=True with stateful=True.
- '
return_state=Truemeans "give me the state tensors as outputs"' - '
stateful=Truemeans "carry state from one batch to the next"'
They solve different problems. If you only want to inspect or reuse the state after a call, return_state is the relevant option.
Common Pitfalls
The most common mistake is expecting all time-step outputs when return_sequences=False, which is the default. In that case Keras only returns the final output.
Another issue is assuming the final state and the output sequence mean the same thing for every recurrent cell type. That is mostly true for SimpleRNN, but not for LSTM.
People also mix up return_state and stateful. One exposes state tensors in the layer output, while the other changes how batches are processed across calls.
Summary
- Use
return_sequences=Trueto get output for every time step. - Use
return_state=Truewhen you also need the final state tensors explicitly. - For
SimpleRNN, the last output and final state often match. - For
LSTM, the final hidden state and final cell state are separate values. - Decide first whether you need the full sequence, the final state, or both.
Related reading
- TensorFlow getting all states from a \`RNN\`
- Tensorflow GPU Could not load dynamic library 'cusolver64_10.dll'; dlerror cusolver64_10.dll not found
- TensorFlow GPU is cudnn optional? Couldn't open CUDA library libcudnn.so
- Tensorflow GPU utilization only 60 GTX 1070
- TensorFlow getting elements of every row for specific columns
- Tensorflow Getting scalar tensor value as int for pass to set_shape
- TensorFlow getting variable by name
- Tensorflow GradientTape Gradients does not exist for variables intermittently
.png&w=3840&q=75)
Tackling System Design Interview Problems
A short course that equips you with the skills to approach system design interviews methodically.
Start the free courseTrack what you have practised
A free account saves your progress, solutions and study plan across every problem on Codemia.
ML System Design practice on Codemia
Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.