Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

LSTM Gate Equations › Comparing with PyTorch

Michael BrenndoerferUpdated May 11, 20252 min read

Language AI HandbookAll chapters

Part 11 of 67Recurrent Neural Networks

Chapter 78LSTM Gate Equations › Comparing with PyTorch

Reading section 12 of 17 · Comparing with PyTorch · 2 minSection 12 of 17 · 2 min

Section 12 of 17 · 2 min read

LSTM Gate Equations › Comparing with PyTorchOne section · Sections in this chapter12 / 17 ↓

Sections in this chapter17

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Let's verify our implementation matches PyTorch's LSTM by comparing outputs:

In[19]:

Code

import torch
import torch.nn as nn

# Set seeds for reproducibility
np.random.seed(42)
torch.manual_seed(42)

# Dimensions
input_dim = 8
hidden_dim = 16
seq_len = 5

# Create our LSTM
our_lstm = LSTMCell(input_dim, hidden_dim)

# Create PyTorch LSTM and copy weights
torch_lstm = nn.LSTMCell(input_dim, hidden_dim)

# PyTorch stores weights concatenated: [W_i, W_f, W_c, W_o]
# We need to reorganize our weights to match
with torch.no_grad():
    # Input-hidden weights (ih): concatenated [i, f, c, o]
    torch_lstm.weight_ih.copy_(
        torch.tensor(
            np.vstack([our_lstm.W_i, our_lstm.W_f, our_lstm.W_c, our_lstm.W_o]),
            dtype=torch.float32,
        )
    )

    # Hidden-hidden weights (hh): concatenated [i, f, c, o]
    torch_lstm.weight_hh.copy_(
        torch.tensor(
            np.vstack([our_lstm.U_i, our_lstm.U_f, our_lstm.U_c, our_lstm.U_o]),
            dtype=torch.float32,
        )
    )

    # Biases
    torch_lstm.bias_ih.copy_(
        torch.tensor(
            np.concatenate(
                [our_lstm.b_i, our_lstm.b_f, our_lstm.b_c, our_lstm.b_o]
            ),
            dtype=torch.float32,
        )
    )

    torch_lstm.bias_hh.copy_(torch.zeros(4 * hidden_dim))

# Test input
x = np.random.randn(input_dim).astype(np.float32)
h_prev = np.zeros(hidden_dim, dtype=np.float32)
c_prev = np.zeros(hidden_dim, dtype=np.float32)

# Our forward pass
our_h, our_c = our_lstm.forward(x, h_prev, c_prev)

# PyTorch forward pass
x_torch = torch.tensor(x).unsqueeze(0)  # Add batch dimension
h_torch = torch.tensor(h_prev).unsqueeze(0)
c_torch = torch.tensor(c_prev).unsqueeze(0)
torch_h, torch_c = torch_lstm(x_torch, (h_torch, c_torch))

Out[20]:

Console

Comparison between our LSTM and PyTorch LSTM:

Max absolute difference in hidden state: 3.22e-08
Max absolute difference in cell state: 6.41e-08

✓ Implementations match!

Our h_t[:5]: [ 0.0283901  -0.36721011 -0.10207159  0.0654979   0.1483921 ]
PyTorch h[:5]: [ 0.02839009 -0.36721012 -0.10207159  0.06549793  0.1483921 ]

The outputs match to numerical precision, confirming our implementation is correct. The small differences (on the order of $10^{-7}$ ) come from floating-point arithmetic variations between NumPy and PyTorch.