TensorFlow
tf.nn.rnn
`RNN`
neural networks
machine learning

What is the equivalent of tf.nn.rnn in new versions of 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.

Practice ML system design

Introduction

Older TensorFlow 1.x code often used tf.nn.rnn or the closely related static_rnn and dynamic_rnn helpers to run a recurrent cell over a sequence. In TensorFlow 2.x, the practical replacement is the Keras recurrent stack: tf.keras.layers.RNN when you want direct control over the cell, or specialized layers such as LSTM, GRU, and SimpleRNN for most day-to-day work.

What Replaced tf.nn.rnn

The old API exposed the unrolling mechanics more directly. That was useful for low-level graph construction, but it also made routine model code noisy. TensorFlow 2.x shifted sequence modeling toward Keras layers that work with eager execution, Model.fit, and SavedModel export.

Think about the migration like this:

  • If old code created a cell and passed it into an RNN helper, the closest conceptual replacement is tf.keras.layers.RNN(cell).
  • If old code just needed a standard recurrent layer, use tf.keras.layers.SimpleRNN, tf.keras.layers.GRU, or tf.keras.layers.LSTM.
  • If you are only trying to keep legacy code alive, tf.compat.v1 still exists, but it is a migration bridge, not the long-term API.

The Keras form also makes batching and masking more obvious. Inputs are normally shaped as batch, timesteps, features, and options such as return_sequences and return_state cover most use cases that previously required lower-level wiring.

Using Keras Recurrent Layers

For standard sequence classification or forecasting, a dedicated recurrent layer is the simplest choice. The example below creates a tiny classifier on synthetic sequence data. It is self-contained and can be run as-is in TensorFlow 2.x.

python
1import numpy as np
2import tensorflow as tf
3
4# 256 examples, each with 12 time steps and 4 features.
5X = np.random.randn(256, 12, 4).astype("float32")
6y = (X.mean(axis=(1, 2)) > 0).astype("float32")
7
8model = tf.keras.Sequential(
9    [
10        tf.keras.layers.Input(shape=(12, 4)),
11        tf.keras.layers.GRU(16),
12        tf.keras.layers.Dense(1, activation="sigmoid"),
13    ]
14)
15
16model.compile(optimizer="adam", loss="binary_crossentropy", metrics=["accuracy"])
17model.fit(X, y, epochs=3, batch_size=32, verbose=0)
18
19predictions = model.predict(X[:3], verbose=0)
20print(predictions)

This is the pattern most teams should prefer. You get optimized kernels, simple checkpointing, and cleaner integration with the rest of the TensorFlow 2 stack.

When tf.keras.layers.RNN Is the Better Match

Sometimes the old tf.nn.rnn usage was not about a built-in architecture. It was about driving a custom cell. In that case, the direct replacement is the generic RNN layer around a cell object that defines state_size, output_size, and call.

python
1import tensorflow as tf
2
3class AdditiveCell(tf.keras.layers.Layer):
4    def __init__(self, units):
5        super().__init__()
6        self.units = units
7        self.state_size = units
8        self.output_size = units
9
10    def build(self, input_shape):
11        feature_count = input_shape[-1]
12        self.kernel = self.add_weight(shape=(feature_count, self.units))
13        self.recurrent_kernel = self.add_weight(shape=(self.units, self.units))
14        self.bias = self.add_weight(shape=(self.units,), initializer="zeros")
15
16    def call(self, inputs, states):
17        prev = states[0]
18        output = tf.tanh(
19            tf.matmul(inputs, self.kernel)
20            + tf.matmul(prev, self.recurrent_kernel)
21            + self.bias
22        )
23        return output, [output]
24
25layer = tf.keras.layers.RNN(AdditiveCell(8), return_sequences=True)
26sequence = tf.random.normal((2, 5, 3))
27result = layer(sequence)
28print(result.shape)

That structure is much closer to the old cell-based workflow, but it still uses the modern Keras execution model.

Migration Notes for Legacy Code

A few old patterns need explicit translation:

  • 'sequence_length behavior is often replaced with masking. An embedding layer with mask_zero=True or a manual mask tensor is the modern route.'
  • Manual loop control around graph sessions disappears in TensorFlow 2.x because eager execution is the default.
  • If old code returned both the full output sequence and the last state, use return_sequences=True and return_state=True.
  • If the model depended on placeholders and feed dictionaries, rewrite the input path using NumPy arrays or tf.data.Dataset.

The migration is usually easier if you first reproduce the old tensor shapes, then replace pieces one layer at a time instead of rewriting the entire model in one jump.

Common Pitfalls

The most common mistake is assuming there is a one-line rename for tf.nn.rnn. There is not. The replacement depends on what the old code was doing.

Another frequent issue is passing data with the wrong shape. Keras recurrent layers expect three dimensions: batch size, time steps, and features. A two-dimensional matrix will trigger shape errors or silently model the wrong structure.

A third pitfall is overusing tf.compat.v1. It can unblock a migration, but code that stays there misses the main benefits of TensorFlow 2.x, including simpler debugging and better Keras integration.

Summary

  • The modern replacement for most tf.nn.rnn code is a Keras recurrent layer.
  • Use SimpleRNN, GRU, or LSTM for standard models.
  • Use tf.keras.layers.RNN(custom_cell) when the old code depended on a custom cell.
  • Translate old shape handling, masking, and state returns explicitly during migration.
  • Treat tf.compat.v1 as temporary compatibility code, not the target design.

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.