Keras
LSTM
input dimension
machine learning
neural networks

Keras LSTM input dimension setting

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

LSTM layers in Keras expect sequence data, so their inputs have one more dimension than a standard dense network. Most confusion comes from mixing up samples, timesteps, and features. Once those three roles are clear, the input shape rules become much easier to apply.

The Core Shape Rule

An LSTM input has this conceptual structure:

text
(samples, timesteps, features)
  • 'samples is the number of training examples in the batch or dataset'
  • 'timesteps is the sequence length for one example'
  • 'features is the number of values available at each timestep'

When you define the model, Keras usually asks only for input_shape=(timesteps, features). The sample dimension stays dynamic.

python
1import tensorflow as tf
2
3model = tf.keras.Sequential([
4    tf.keras.layers.Input(shape=(12, 5)),
5    tf.keras.layers.LSTM(32),
6    tf.keras.layers.Dense(1),
7])
8
9model.summary()

This model expects training data with shape (number_of_samples, 12, 5).

Reading the Dimensions Correctly

Suppose you have 100 examples. Each example is a sequence of 12 steps, and each step has 5 measurements. Then the correct shape is:

python
(100, 12, 5)

A runnable example:

python
1import numpy as np
2import tensorflow as tf
3
4x_train = np.random.rand(100, 12, 5).astype("float32")
5y_train = np.random.rand(100, 1).astype("float32")
6
7model = tf.keras.Sequential([
8    tf.keras.layers.Input(shape=(12, 5)),
9    tf.keras.layers.LSTM(16),
10    tf.keras.layers.Dense(1),
11])
12
13model.compile(optimizer="adam", loss="mse")
14model.fit(x_train, y_train, epochs=2, batch_size=8, verbose=0)

The key point is that x_train is rank 3. If you pass rank 2 data, Keras cannot tell which axis is time.

Univariate Time Series

A very common mistake happens with one-dimensional sequences. If each example has 20 timesteps and one value per timestep, the input shape is (20, 1), not (20,).

python
1import numpy as np
2import tensorflow as tf
3
4x_train = np.random.rand(200, 20, 1).astype("float32")
5y_train = np.random.rand(200, 1).astype("float32")
6
7model = tf.keras.Sequential([
8    tf.keras.layers.Input(shape=(20, 1)),
9    tf.keras.layers.LSTM(32),
10    tf.keras.layers.Dense(1),
11])
12
13model.compile(optimizer="adam", loss="mse")
14model.fit(x_train, y_train, epochs=2, batch_size=16, verbose=0)

That trailing 1 matters because the LSTM still expects a features axis even when there is only one feature.

Stacked LSTMs and return_sequences

If one LSTM feeds another LSTM, the first one usually must return the full sequence:

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

Without return_sequences=True, the first LSTM returns only its final output vector. The second LSTM then receives rank 2 data instead of a sequence and the shape no longer matches.

When batch_input_shape Is Needed

Most models should use input_shape, not batch_input_shape. The batch size only needs to be fixed in special cases such as stateful LSTMs.

python
1import tensorflow as tf
2
3model = tf.keras.Sequential([
4    tf.keras.layers.LSTM(
5        8,
6        stateful=True,
7        batch_input_shape=(4, 10, 3),
8    ),
9    tf.keras.layers.Dense(1),
10])

Here the batch size is fixed at 4. Stateful recurrent models are more restrictive, so use them only when you really need state to persist across batches.

Common Pitfalls

The most common error is swapping timesteps and features. Data shaped like (samples, features, timesteps) is not the same thing as (samples, timesteps, features).

Another mistake is forgetting the features axis for univariate time series. A sequence of length 20 with one value at each step still needs shape (20, 1) at the layer level.

People also forget return_sequences=True when stacking recurrent layers. That causes later LSTM layers to receive the wrong rank.

Finally, many shape bugs are actually data-preparation bugs. Always print your training array shapes before blaming the model definition.

Summary

  • Keras LSTMs expect input shaped like (samples, timesteps, features).
  • In model definitions, use input_shape=(timesteps, features) for most cases.
  • A single-feature sequence still needs a features axis, such as (20, 1).
  • Stacked LSTMs usually require return_sequences=True on earlier recurrent layers.
  • Use batch_input_shape only for special cases such as stateful models.

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.