TensorFlow
dynamic_rnn
`RNN`
machine learning
neural networks

Get last output of dynamic_rnn in tensorflow?

Master System Design with Codemia

Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.

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].

python
1import tensorflow as tf
2
3tf.compat.v1.disable_eager_execution()
4
5inputs = tf.compat.v1.placeholder(tf.float32, [None, 5, 3])
6cell = tf.compat.v1.nn.rnn_cell.GRUCell(4)
7
8outputs, state = tf.compat.v1.nn.dynamic_rnn(
9    cell,
10    inputs,
11    dtype=tf.float32,
12)
13
14print(outputs)
15print(state)

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.

python
last_output = outputs[:, -1, :]
print(last_output)

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.

python
1import tensorflow as tf
2
3tf.compat.v1.disable_eager_execution()
4
5inputs = tf.compat.v1.placeholder(tf.float32, [None, 5, 3])
6lengths = tf.compat.v1.placeholder(tf.int32, [None])
7cell = tf.compat.v1.nn.rnn_cell.GRUCell(4)
8
9outputs, state = tf.compat.v1.nn.dynamic_rnn(
10    cell,
11    inputs,
12    sequence_length=lengths,
13    dtype=tf.float32,
14)
15
16batch_size = tf.shape(outputs)[0]
17max_time = tf.shape(outputs)[1]
18output_size = tf.shape(outputs)[2]
19
20flat = tf.reshape(outputs, [-1, output_size])
21index = tf.range(batch_size) * max_time + (lengths - 1)
22last_valid_output = tf.gather(flat, index)
23
24print(last_valid_output)

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:

python
1gru_cell = tf.compat.v1.nn.rnn_cell.GRUCell(4)
2outputs, state = tf.compat.v1.nn.dynamic_rnn(gru_cell, inputs, dtype=tf.float32)
3
4final_representation = state

For an LSTM, the final state contains both cell state and hidden state. In most classifier-style uses, you want the hidden state h.

python
1lstm_cell = tf.compat.v1.nn.rnn_cell.LSTMCell(4)
2outputs, state = tf.compat.v1.nn.dynamic_rnn(lstm_cell, inputs, dtype=tf.float32)
3
4last_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_rnn returns both per-step outputs and a final state.'
  • For fixed-length sequences, outputs[:, -1, :] is usually enough.
  • For padded variable-length data, use sequence_length and gather the last valid step.
  • For GRU and similar cells, state is often the representation you actually want.
  • For LSTM, the final hidden state is usually state.h.

Course illustration
Course illustration

All Rights Reserved.