PyTorch LSTM: shapes, hidden states, and sequence output

PyTorch’s nn.LSTM processes a sequence and returns two things: an output for each time step from the last layer, and a pair of final states, (h_n, c_n). With batch_first=True, the batched input uses (batch, time, features). The final states keep a different layout, with layers and directions first.

This guide uses five runnable CPU examples to distinguish those tensors. We executed 37 checks with Python 3.12.14 and PyTorch 2.15.0+cpu, including a separate scalar calculation, variable-length sequences, and state carried between chunks. Each code block runs independently with PyTorch installed. These are correctness examples, not evidence of forecasting accuracy or GPU performance.

Start with one hidden unit

An LSTM maintains a hidden state and a cell state. In this first example, we choose the weights so that the calculation can be checked by hand. All weights and biases are zero except the input weight for the candidate cell value. The LSTM API reference specifies the gate order used to locate that coefficient.

import torch

x = torch.tensor([[[1.0], [2.0], [3.0]]])
layer = torch.nn.LSTM(1, 1, batch_first=True)
with torch.no_grad():
    for parameter in layer.parameters():
        parameter.zero_()
    # Gate rows are input, forget, candidate, output.
    layer.weight_ih_l0[2, 0] = 1.0
output, (h_n, c_n) = layer(x)
print('Output shape:', tuple(output.shape))
print('Hidden values:', [round(v, 6) for v in output.flatten().tolist()])
print('Final hidden:', round(h_n.item(), 6))
print('Final cell:', round(c_n.item(), 6))

Observed output:

Output shape: (1, 3, 1)
Hidden values: [0.1817, 0.293282, 0.341238]
Final hidden: 0.341238
Final cell: 0.833734

With these coefficients, the input, forget, and output gates each equal 0.5. Starting from a zero cell, our independent calculation is c_t = 0.5 × c_(t-1) + 0.5 × tanh(x_t), followed by h_t = 0.5 × tanh(c_t). For the first input, 1, the cell becomes approximately 0.380797 and the hidden value becomes 0.181700.

After all three inputs, the cell is 0.833734 and the hidden state is 0.341238. They are different quantities. The last entry of this one-direction output matches the final hidden state, not the final cell. Setting the weights this way exposes the recurrence; it does not train a useful sequence model.

Separate batch layout from state layout

Suppose each example contains four time steps with three measured features per step. A batch of two examples has shape (2, 4, 3). Here, input_size=3 describes features at one step, not the batch size or sequence length.

We create two layers with identical parameters and change only their input layout. Transposing the batch and time axes makes their results comparable:

import torch

torch.manual_seed(7)
x = torch.arange(24, dtype=torch.float32).reshape(2, 4, 3) / 10
batch_layer = torch.nn.LSTM(3, 5, num_layers=2, batch_first=True)
time_layer = torch.nn.LSTM(3, 5, num_layers=2)
time_layer.load_state_dict(batch_layer.state_dict())
batch_out, (batch_h, batch_c) = batch_layer(x)
time_out, (time_h, time_c) = time_layer(x.transpose(0, 1))
print('Batch-first output:', tuple(batch_out.shape))
print('Time-first output:', tuple(time_out.shape))
print('Hidden and cell:', tuple(batch_h.shape), tuple(batch_c.shape))
print('Same output:', torch.allclose(batch_out, time_out.transpose(0, 1)))
print('Same final states:', torch.allclose(batch_h, time_h)
      and torch.allclose(batch_c, time_c))

Observed output:

Batch-first output: (2, 4, 5)
Time-first output: (4, 2, 5)
Hidden and cell: (2, 2, 5) (2, 2, 5)
Same output: True
Same final states: True

The batch-first output is (2, 4, 5): two sequences, four positions, and five hidden features. The time-first output is (4, 2, 5). After transposing it back, every value matches within the comparison tolerance.

Both final-state tensors are (2, 2, 5) here: two layers, two batch members, and hidden width five. Changing batch_first does not move the batch axis of those states. Our separate checks also confirmed that omitting the initial states matches passing explicit zero tensors, and that an incorrect initial layer count raises an error.

Use an axis transpose rather than a reshape when converting time-first and batch-first data. Our reshape counterexample has the desired dimensions but different element placement. Check the batches your DataLoader actually produces before adding a conversion.

Read both directions from the final states

The next example has two layers, two directions, and hidden width three. Its output feature width is six because each position contains forward and reverse features. The first state axis has four entries: two directions for each of two layers.

import torch

torch.manual_seed(7)
x = torch.arange(16, dtype=torch.float32).reshape(2, 4, 2) / 10
layer = torch.nn.LSTM(2, 3, num_layers=2, bidirectional=True, batch_first=True)
output, (h_n, c_n) = layer(x)
# Separate layer, direction, batch, hidden width.
states = h_n.reshape(2, 2, 2, 3)
summary = torch.cat((states[-1, 0], states[-1, 1]), dim=-1)
print('Output:', tuple(output.shape))
print('Hidden and cell:', tuple(h_n.shape), tuple(c_n.shape))
print('Summary:', tuple(summary.shape))
print('Forward endpoint:', torch.allclose(output[:, -1, :3], states[-1, 0]))
print('Reverse endpoint:', torch.allclose(output[:, 0, 3:], states[-1, 1]))
print('Last output is full summary:', torch.allclose(output[:, -1], summary))

Observed output:

Output: (2, 4, 6)
Hidden and cell: (4, 2, 3) (4, 2, 3)
Summary: (2, 6)
Forward endpoint: True
Reverse endpoint: True
Last output is full summary: False

Reshaping h_n into explicit layer and direction axes lets us select the last layer and concatenate its two final states. The result has shape (2, 6), one six-value summary per batch member. This reshape splits a combined axis; it is not the batch/time rearrangement warned about above.

The forward final state matches the forward half of the output at the last position. The reverse final state matches the reverse half at the first position. The final output position alone fails our comparison with the full summary. That is why output[:, -1] is the wrong shortcut for this bidirectional example.

A model that needs future sequence elements cannot provide the same bidirectional result before those elements arrive. Choose the direction setting around the information available when you will make the prediction. For the wider choice between recurrent models and other designs, see our neural network architecture guide.

Keep padding out of the recurrent calculation

Our next batch contains sequences of lengths two and four. The shorter one is padded with zeros to form a rectangular tensor. Passing that tensor directly through the LSTM performs four steps for both examples; a zero input is still a step in the recurrence.

pack_padded_sequence lets us specify valid lengths. The lengths tensor stays on CPU. We set enforce_sorted=False because the lengths are not in descending order. Each parameter is fixed to 0.1 so the comparison is repeatable.

import torch

from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence

x = torch.tensor([[[1.], [2.], [0.], [0.]],
                  [[3.], [4.], [5.], [6.]]])
lengths = torch.tensor([2, 4], dtype=torch.long)  # CPU tensor
layer = torch.nn.LSTM(1, 2, batch_first=True)
with torch.no_grad():
    for parameter in layer.parameters():
        parameter.fill_(0.1)
packed = pack_padded_sequence(x, lengths, batch_first=True, enforce_sorted=False)
packed_out, (packed_h, packed_c) = layer(packed)
padded_out, recovered = pad_packed_sequence(packed_out, batch_first=True, total_length=4)
_, (single_h, single_c) = layer(x[:1, :2])
_, (plain_h, plain_c) = layer(x)
print('Recovered lengths:', recovered.tolist())
print('Padded output:', tuple(padded_out.shape))
print('Short sequence matches alone:', torch.allclose(packed_h[:, :1], single_h))
print('Plain padding matches:', torch.allclose(plain_h[:, :1], single_h))
print('Unpacked tail is zero:', torch.count_nonzero(padded_out[0, 2:]).item() == 0)

Observed output:

Recovered lengths: [2, 4]
Padded output: (2, 4, 2)
Short sequence matches alone: True
Plain padding matches: False
Unpacked tail is zero: True

The packed hidden state for the shorter sequence matches processing its two real steps alone. Our additional check found the same match for its cell state. Plain zero padding does not produce that final hidden state: the two extra steps changed it.

We also checked the longer sequence separately and compared every valid output value for both examples. The returned state preserves the original batch order even though packing sorted the sequences internally. This matters when matching representations back to labels or identifiers.

pad_packed_sequence restores a rectangular output and returns the lengths. In our result, the shorter example’s unused tail is zero, and total_length=4 gives the expected time dimension. Those zeros are padding in the unpacked result, not two additional recurrent computations.

Lengths describe valid leading steps. Do not use an arbitrary mask with gaps in the middle as though it were a length. Our checks rejected a zero length and rejected unsorted lengths under the default sorting requirement. Packing inputs also does not automatically exclude padded positions from a loss you compute after unpacking; choose the valid positions deliberately.

Carry state only between chunks of the same sequence

For a one-direction LSTM with unchanged weights and no dropout, we split one sequence into two chunks. Passing both states from the first chunk into the second reproduces the full-sequence result in this test:

import torch

torch.manual_seed(7)
x = torch.arange(12, dtype=torch.float32).reshape(1, 6, 2) / 10
layer = torch.nn.LSTM(2, 3, batch_first=True)
whole, whole_state = layer(x)
first, state = layer(x[:, :3])
second, final_state = layer(x[:, 3:], state)
reset_second, _ = layer(x[:, 3:])
print('Carried output matches:', torch.allclose(torch.cat((first, second), dim=1), whole))
print('Carried states match:', all(torch.allclose(a, b) for a, b in zip(final_state, whole_state)))
print('Reset output matches:', torch.allclose(reset_second, whole[:, 3:]))

Observed output:

Carried output matches: True
Carried states match: True
Reset output matches: False

Resetting the state for the second chunk changes its output. Carrying the state preserves the connection across the boundary. Batch positions must still refer to the same sequences; do not carry state blindly between shuffled batches or unrelated examples.

For truncated training histories, tuple(s.detach() for s in state) preserves the state values while disconnecting their earlier computation graph. We checked both properties locally. Detaching changes which earlier steps can receive gradients; it does not reset the stored values or prove that a particular training strategy will converge.

This chunk comparison does not cover bidirectional processing, dropout, or parameter updates between chunks. If you need to resume a training run, distinguish these sequence states from the full model and optimizer checkpoint. Our save-and-load guide covers checkpoint recovery.

Check these details before training

• Feature width: the final input dimension must match input_size. Our mismatched-width test raised an error.
• State axes: write down layers, directions, batch, and hidden width before indexing h_n.
• Sequence boundaries: use true lengths or deliberate state resets; padding and unrelated batches are not continuations.
• Output meaning: an LSTM representation is not yet a class probability or forecast. Define the prediction head and target separately.

All examples leave projections disabled and dropout at zero. The state widths shown here should not be copied into a projected LSTM without checking its configuration. We also verified an unbatched input against a batch of one and checked finite input and parameter gradients in a backward pass.

The tests establish the displayed values, layouts, and state behavior. They do not evaluate a dataset, task accuracy, CUDA, mixed precision, throughput, or long-run training stability. The environment emitted a missing-NumPy initialization warning; these tensor-only checks completed without NumPy.