tf.keras.layers.LSTM() is a recurrent layer that processes sequential data step by step, maintaining an internal memory state that lets it capture long-range dependencies.

Syntax

tf.keras.layers.LSTM(units, return_sequences=False)

Deep Dive Course

An LSTM, Long Short-Term Memory, layer processes a sequence one timestep at a time, maintaining both a hidden state and a separate cell state that acts as a longer-term memory, using learned gates to decide what information to keep, forget, or output at each step. This gating mechanism specifically helps LSTMs retain information across many timesteps, addressing the vanishing-gradient problem that made simpler recurrent layers struggle to learn long-range dependencies. By default, return_sequences=False, so the layer only outputs its final hidden state after processing the whole sequence; setting it to True returns the hidden state at every timestep instead, needed when stacking multiple recurrent layers.

1Understanding tf.keras.layers.LSTM()

💡

Set return_sequences=True whenever you're stacking one LSTM layer on top of another, since each layer after the first needs a full sequence of hidden states as its input, not just the single final one that return_sequences=False produces.

editor.html

import tensorflow as tf
from tensorflow.keras import layers

layer = layers.LSTM(16)
x = tf.zeros([4, 10, 8])
output = layer(x)
print(output.shape)

localhost:3000

2Practical Example

Here is a real-world application of tf.keras.layers.LSTM() showing how it is used in production TensorFlow code.

editor.html

import tensorflow as tf
from tensorflow.keras import layers

layer = layers.LSTM(16, return_sequences=True)
x = tf.zeros([4, 10, 8])
output = layer(x)
print(output.shape)

localhost:3000

3Best Practices

Follow these guidelines when working with tf.keras.layers.LSTM():

1. Set return_sequences=True on every LSTM layer except the last one when stacking multiple recurrent layers, since intermediate layers need a full sequence as output

2. Shape input data as (batch_size, timesteps, features) before passing it to an LSTM layer, matching its expected 3D input format

3. Consider a simpler GRU layer as a faster alternative to LSTM when maximum long-range memory capacity isn't essential, since GRU uses fewer gates and parameters

⚠️

Tip: Set return_sequences=True whenever you're stacking one LSTM layer on top of another, since each layer after the first needs a full sequence of hidden states as its input, not just the single final one that return_sequences=False produces.

editor.html

import tensorflow as tf
from tensorflow.keras import layers

layer = layers.LSTM(16)
x = tf.zeros([4, 10, 8])
output = layer(x)
print(output.shape)

localhost:3000

Examples

Example 01Basic Usage

import tensorflow as tf
from tensorflow.keras import layers

layer = layers.LSTM(16)
x = tf.zeros([4, 10, 8])
output = layer(x)
print(output.shape)

Example 02Advanced Example

import tensorflow as tf
from tensorflow.keras import layers

layer = layers.LSTM(16, return_sequences=True)
x = tf.zeros([4, 10, 8])
output = layer(x)
print(output.shape)

Best Practices

Set return_sequences=True on every LSTM layer except the last one when stacking multiple recurrent layers, since intermediate layers need a full sequence as output
Shape input data as (batch_size, timesteps, features) before passing it to an LSTM layer, matching its expected 3D input format
Consider a simpler GRU layer as a faster alternative to LSTM when maximum long-range memory capacity isn't essential, since GRU uses fewer gates and parameters

Interview Question

Why do LSTM layers generally handle long sequences better than a simpler recurrent layer, like SimpleRNN?▼

Hint: Think about the vanishing gradient problem and what mechanism LSTM specifically adds to address it.

A simple recurrent layer repeatedly multiplies its hidden state by the same weight matrix at every timestep, and during backpropagation through time, gradients get multiplied by that same matrix repeatedly too — over many timesteps, this tends to shrink gradients toward zero, the vanishing gradient problem, making it very difficult for the network to learn dependencies between distant timesteps, since the training signal from a late error barely reaches back to influence early inputs. An LSTM addresses this with its separate cell state and learned gating mechanism, specifically a forget gate that can be trained to keep information flowing through the cell state largely unchanged across many timesteps when needed, providing something closer to a direct path for gradients to flow backward without repeatedly shrinking — this gated memory mechanism is exactly what lets LSTMs learn dependencies across much longer sequences than a simple recurrent layer typically can.

Exercises

MediumPractice using tf.keras.layers.LSTM() in a real scenario.

View Solution

import tensorflow as tf
from tensorflow.keras import layers

layer = layers.LSTM(16)
x = tf.zeros([4, 10, 8])
output = layer(x)
print(output.shape)

Frequently Asked Questions

Why do LSTM layers generally handle long sequences better than a simpler recurrent layer, like SimpleRNN?▼

Related Functions

tf-keras-layers-embeddingtf-keras-sequentialmodel-fit