LSTM
Keras
multidimensional input
deep learning
neural networks

Multi dimensional input for LSTM in Keras

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

Keras LSTM layers expect sequence data, not a flat feature table. The key idea is that each training example is a sequence of timesteps, and each timestep contains one or more features, so the standard input shape is three-dimensional.

The Shape Keras Expects

For an LSTM in Keras or tf.keras, the input tensor shape is:

  • 'samples'
  • 'timesteps'
  • 'features'

When you define the layer, you specify only timesteps and features, because the batch dimension is handled automatically.

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

That model expects each sample to contain 10 time steps and 5 numeric features at each step.

What "Multidimensional" Usually Means

In most practical questions, "multidimensional input" means one of two things.

The first meaning is simple multivariate sequence data, such as temperature, pressure, and humidity values observed at each time step. That fits directly into the normal LSTM shape of (samples, timesteps, features).

The second meaning is more complex structured data at each step, such as an image frame, a grid, or a vector of embeddings. In those cases, you often need an additional preprocessing layer before the LSTM or a different recurrent architecture.

A Runnable Multivariate Example

This example trains an LSTM on synthetic data with 100 samples, 12 time steps, and 3 features per step.

python
1import numpy as np
2import tensorflow as tf
3
4np.random.seed(0)
5
6a = np.random.rand(100, 12, 3).astype("float32")
7y = a[:, :, 0].sum(axis=1) > 6.0
8y = y.astype("float32")
9
10model = tf.keras.Sequential([
11    tf.keras.layers.Input(shape=(12, 3)),
12    tf.keras.layers.LSTM(16),
13    tf.keras.layers.Dense(1, activation="sigmoid")
14])
15
16model.compile(optimizer="adam", loss="binary_crossentropy", metrics=["accuracy"])
17model.fit(a, y, epochs=5, batch_size=16, verbose=0)
18
19print(model.predict(a[:2], verbose=0))

The important detail is not the toy target. It is the shape of a, which is (100, 12, 3).

Stacking LSTM Layers

If you want multiple LSTM layers, the first recurrent layer must return the full sequence, not just the final output. That is what return_sequences=True does.

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

Without return_sequences=True on the first layer, the second LSTM would receive a two-dimensional tensor instead of a sequence.

What If Each Timestep Is Itself High-Dimensional?

Sometimes each timestep is not just a feature vector. It may be an image, spectrogram slice, or other structured tensor.

For example, if each timestep is a 32 x 32 image with 3 channels, the raw input shape would conceptually be (samples, timesteps, height, width, channels). A plain LSTM does not consume that directly.

Common solutions are:

  • extract per-frame features with TimeDistributed and a smaller network
  • use ConvLSTM2D if the sequence is genuinely spatiotemporal
  • reshape only if the flattening still preserves the information you care about

A TimeDistributed example looks like this.

python
1import tensorflow as tf
2
3model = tf.keras.Sequential([
4    tf.keras.layers.Input(shape=(10, 28, 28, 1)),
5    tf.keras.layers.TimeDistributed(tf.keras.layers.Flatten()),
6    tf.keras.layers.LSTM(32),
7    tf.keras.layers.Dense(10, activation="softmax")
8])

That converts each frame into a vector before the LSTM processes the temporal sequence.

Sequence Length and Padding

Real datasets often contain sequences of unequal length. Keras can still handle them if you pad shorter sequences to a common length and optionally use masking.

python
1import tensorflow as tf
2
3model = tf.keras.Sequential([
4    tf.keras.layers.Masking(mask_value=0.0, input_shape=(15, 4)),
5    tf.keras.layers.LSTM(16),
6    tf.keras.layers.Dense(1)
7])

This is useful when some timesteps are just placeholders added to make the batch rectangular.

Common Pitfalls

The most common mistake is swapping the timesteps and features axes. If your model expects (12, 3) and you pass (3, 12), training may still run but the sequence meaning is wrong.

Another mistake is flattening the whole sample before the LSTM. That destroys the temporal structure the recurrent layer is supposed to learn.

A third issue is forgetting return_sequences=True when stacking recurrent layers.

Finally, if your per-timestep data is image-like or grid-like, a plain LSTM may be the wrong layer. Consider TimeDistributed, CNN feature extraction, or ConvLSTM2D instead.

Summary

  • A standard Keras LSTM expects input shaped as (samples, timesteps, features).
  • Multivariate sequence data fits naturally into that three-dimensional structure.
  • Use return_sequences=True when feeding one LSTM into another.
  • Use masking or padding when sequence lengths vary.
  • For image-like data per timestep, preprocess with TimeDistributed or use ConvLSTM2D.
  • Most shape errors come from mixing up axes or flattening away the time dimension.

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.