TensorFlow
`RNN`
neural networks
machine learning
artificial intelligence

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.

Practice ML system design

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.

python
1import tensorflow as tf
2
3x = tf.random.normal((2, 5, 3))
4
5layer = tf.keras.layers.SimpleRNN(
6    4,
7    return_sequences=True
8)
9
10all_outputs = layer(x)
11
12print(all_outputs.shape)  # (2, 5, 4)

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:

python
1import tensorflow as tf
2
3x = tf.random.normal((2, 5, 3))
4
5layer = tf.keras.layers.SimpleRNN(
6    4,
7    return_sequences=True,
8    return_state=True
9)
10
11all_outputs, final_state = layer(x)
12
13print(all_outputs.shape)  # (2, 5, 4)
14print(final_state.shape)  # (2, 4)

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.

python
1import tensorflow as tf
2
3x = tf.random.normal((2, 5, 3))
4
5lstm = tf.keras.layers.LSTM(
6    8,
7    return_sequences=True,
8    return_state=True
9)
10
11sequence_outputs, final_hidden, final_cell = lstm(x)
12
13print(sequence_outputs.shape)  # (2, 5, 8)
14print(final_hidden.shape)      # (2, 8)
15print(final_cell.shape)        # (2, 8)

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.

python
1import tensorflow as tf
2
3model = tf.keras.Sequential([
4    tf.keras.layers.SimpleRNN(16, return_sequences=True, input_shape=(10, 6)),
5    tf.keras.layers.SimpleRNN(8),
6    tf.keras.layers.Dense(1),
7])
8
9model.summary()

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:

python
1import tensorflow as tf
2
3x = tf.random.normal((1, 4, 2))
4layer = tf.keras.layers.GRU(3, return_sequences=True)
5states_over_time = layer(x)
6
7print(states_over_time.numpy())

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=True means "give me the state tensors as outputs"'
  • 'stateful=True means "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=True to get output for every time step.
  • Use return_state=True when 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
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track 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.

Practice ML system design

All Rights Reserved.