Get last output of dynamic_rnn in tensorflow?
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
With tf.compat.v1.nn.dynamic_rnn, "last output" can mean two different things: the final emitted output tensor or the final recurrent state. Those are often similar for simple cells, but they are not always interchangeable, especially with padded variable-length sequences or LSTM state objects. The right choice depends on what your model actually needs.
Understand What dynamic_rnn Returns
dynamic_rnn returns two values:
- '
outputs, which contains an output for every time step' - '
state, which contains the final state after processing the sequence'
For batch-major input, outputs has shape [batch_size, max_time, output_size].
If every sequence really has length five, the last time step in outputs is easy to obtain. If the batch contains padding, that simple slice may be wrong.
Fixed-Length Sequences: Slice the Last Time Step
When every sequence has the same real length, the last output is simply the last slice along the time dimension.
This is the shortest correct answer for fixed-length data. It works because the final time step for every batch element is meaningful rather than padding.
Variable-Length Sequences Need sequence_length
In real NLP and sequence tasks, batches are often padded. If you use outputs[:, -1, :] there, you may read the output for the padded tail rather than the last real token.
The correct approach is to pass sequence_length into dynamic_rnn and then gather the last valid output for each example.
That pattern aligns the result with the true sequence length of each row in the batch.
Sometimes the Final State Is the Better Answer
If your model conceptually wants the final hidden representation, the returned state is often the cleaner API.
For a GRU:
For an LSTM, the final state contains both cell state and hidden state. In most classifier-style uses, you want the hidden state h.
That distinction matters because state.c and state.h serve different roles inside the LSTM.
Legacy API Versus Modern TensorFlow
dynamic_rnn lives under tf.compat.v1, which is a sign that this is a legacy graph-mode pattern. In modern TensorFlow code, you would usually build the model with Keras RNN layers and ask the layer to return sequences or return state explicitly.
Still, plenty of older production code and research codebases use dynamic_rnn, so understanding the legacy behavior remains useful when maintaining or migrating those models.
Common Pitfalls
The most common mistake is using outputs[:, -1, :] on padded batches without providing sequence_length. That returns the last padded step, not the last real step.
Another issue is assuming that state and the last element of outputs are always the same thing. For some cells they are close enough, but LSTM state objects make the distinction explicit.
Developers also often grab the whole LSTM state and forget that they probably wanted state.h.
Finally, it is easy to mix modern eager-execution expectations with this legacy TensorFlow 1 style API. Be clear about which execution model your code is using.
Summary
- '
dynamic_rnnreturns both per-step outputs and a final state.' - For fixed-length sequences,
outputs[:, -1, :]is usually enough. - For padded variable-length data, use
sequence_lengthand gather the last valid step. - For GRU and similar cells,
stateis often the representation you actually want. - For LSTM, the final hidden state is usually
state.h.
Related reading
- Get length of a dataset in Tensorflow
- get the CUDA and CUDNN version on windows with Anaconda installe
- Get the last output of a dynamic_rnn in TensorFlow
- Getting a prediction from an ONNX model in python
- Get learning rate of keras model
- Get the bounding box coordinates in the TensorFlow object detection API tutorial
- Get the diagonal of a matrix in TensorFlow
- Get the label mappings from label encoder
.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.