Recurrent Neural Networks: Memory as a Moving Point in Space

2026

Python PyTorch RNN

Read the sentence "The cat sat on the ___" and you already know the next word is probably mat, rug, or floor. You didn't get that from the last word alone. You got it from everything you'd read so far, carried forward in your head as you went.

A recurrent neural network (RNN) does the same thing with one idea: keep a single vector as a learned, compressed summary of everything read so far, and update it every time a new word arrives. This tutorial builds that idea from the ground up. We'll picture the summary as a point moving through space, watch a real trained model move it, and then write the whole thing in under 100 lines of PyTorch.

It builds on three earlier tutorials. Introduction to Machine Learning covers data as vectors, mappings between spaces, and gradient descent. Neural Networks Demystified covers hidden layers as embedding spaces. The Company Words Keep covers word embeddings. You don't need them to follow along, but each one fills in a piece we'll use here.

Why We Need Memory

A regular feedforward network takes one fixed-size input and produces one output. That works for a house described by three numbers. It doesn't work well for language, for two reasons.

Sentences have different lengths. "Dogs bark" has two words; this sentence has more. A network with exactly ten input slots has no natural way to read either.

Order matters. "Dog bites man" and "man bites dog" contain the same words. If you add up the word vectors, or otherwise ignore order, the two sentences look identical.

The fix is to stop reading the whole sentence at once. Read it the way you do: one word at a time, keeping a running summary. After each word, update the summary using two things, the summary so far and the new word. That summary is called the hidden state, and the network that updates it is recurrent because the same update is applied again and again.

The One Equation

Here is the entire update rule of a basic RNN:

$$h_t = \tanh\left(W_h\, h_{t-1} + W_x\, x_t + b\right)$$

Read it left to right:

$h_t$ is the hidden state after reading word $t$. It's just a vector of numbers.
$h_{t-1}$ is the hidden state before this word: the memory so far.
$x_t$ is the embedding of the new word (a vector, one per word in the vocabulary).
$W_h$ and $W_x$ are weight matrices, and $b$ is a bias vector. All three are learned.
tanh squashes every number into the range −1 to 1.

So the old state and the new word each get transformed by their own matrix, the results are added, and tanh squashes the sum. It is not "old state plus the word's embedding". Both terms are reshaped by learned weights first, and that is where all the learning lives.

To predict the next word, read a score for every word in the vocabulary off the hidden state, then turn the scores into probabilities:

$$\text{scores} = W_o\, h_t + c, \qquad p(\text{next word}) = \text{softmax}(\text{scores})$$

That's the whole model. If you "unroll" it over a sentence, it looks like the same small machine copied once per word, with the hidden state passed from each copy to the next.

An RNN unrolled over four words, with the same weights at every step — The RNN unrolled over "the cat sat on". Every column uses the **same** W_x, W_h, and W_o. The predictions at the top are the real outputs of the model we train below.

The key detail is that every copy shares the same weights. A ten-word sentence doesn't need ten times as many parameters; it just runs the same update ten times. That's how one network handles sentences of any length.

One step, with real numbers

Let's run one step by hand. Throughout this tutorial we'll use a deliberately tiny model so that every number is visible. It's trained on 16 sentences built from this pattern:

the {cat, dog, bird} {sat on, slept on, ran to} the {mat, rug | door, park} .

Sat and slept are followed by on, then mat or rug; ran is followed by to, then door or park. One combination is held out on purpose: the training data never contains "the bird sat". We'll use that later to test generalization. The vocabulary has 14 tokens (counting the period), each word embedding has 2 numbers, and the hidden state has 2 numbers, so we can plot both directly. The training code is at the end; every number below comes from it.

After reading "the", the model's hidden state is $h_1 = (0.09,\ 1.00)$. Now it reads "cat", whose learned embedding is $x_2 = (-1.30,\ 0.67)$. The trained weights, rounded to two decimals, are:

$$W_h = \begin{bmatrix} -2.18 & 0.28 \\ 0.34 & 3.37 \end{bmatrix}, \quad W_x = \begin{bmatrix} 0.83 & 1.01 \\ 1.73 & -1.24 \end{bmatrix}, \quad b = \begin{bmatrix} -0.42 \\ -0.15 \end{bmatrix}$$

The memory term carries the old state forward:

$$W_h h_1 = \begin{bmatrix} -2.18(0.09) + 0.28(1.00) \\ 0.34(0.09) + 3.37(1.00) \end{bmatrix} = \begin{bmatrix} 0.08 \\ 3.40 \end{bmatrix}$$

The input term brings in the new word:

$$W_x x_2 = \begin{bmatrix} 0.83(-1.30) + 1.01(0.67) \\ 1.73(-1.30) - 1.24(0.67) \end{bmatrix} = \begin{bmatrix} -0.40 \\ -3.08 \end{bmatrix}$$

Add them with the bias, then squash:

$$h_2 = \tanh\left(\begin{bmatrix} 0.08 - 0.40 - 0.42 \\ 3.40 - 3.08 - 0.15 \end{bmatrix}\right) = \tanh\left(\begin{bmatrix} -0.74 \\ 0.17 \end{bmatrix}\right) = \begin{bmatrix} -0.63 \\ 0.17 \end{bmatrix}$$

With full-precision weights the model gets $(-0.64,\ 0.16)$; the small gap is rounding.

Notice the tug-of-war in the second number. The memory term contributes +3.40, and the word "cat" contributes −3.08. The new state is a negotiation between where you were and what you just read. From $h_2$ the model predicts ran 0.35, slept 0.35, sat 0.28: after "the cat", a verb comes next.

Memory as a Moving Point

Here's the picture that makes RNNs click. The hidden state is one point. Our model's hidden state has two numbers, so the point lives on a flat square from −1 to 1 in each direction (tanh can't go further). The point starts at the origin. Each word moves it somewhere new. Wherever it lands decides what the model predicts next.

Hidden-state trajectory for the cat sat on the mat, over shaded prediction regions — The real path of the hidden state while reading "the cat sat on the mat". Each shaded area marks the group receiving the largest combined softmax probability at that point, computed from W_o and c. These are visual summaries of the model’s probabilities, not hard boundaries in the model.

Follow the point:

After reading	Hidden state $h$	Model's top predictions
the	(0.09, 1.00)	cat 0.37, dog 0.37, bird 0.25
the cat	(−0.64, 0.16)	ran 0.35, slept 0.35, sat 0.28
the cat sat	(0.04, −0.98)	on 0.99
the cat sat on	(−1.00, −0.73)	the 0.98
the cat sat on the	(0.97, 0.65)	mat 0.49, rug 0.49
the cat sat on the mat	(−1.00, 1.00)	. 0.99

Each step does two jobs at once. It remembers (the new position depends on the old one) and it positions for prediction (the new position tends to give higher probability to the right next word). The model learns both jobs at the same time: where each word should move the point, and what predictions to make from each position.

Three things in this picture are worth noticing:

The predictions are honest about uncertainty. After "the cat sat on the", the model splits evenly between mat and rug, because in the training data both are equally common. A good language model shouldn't be certain when the data isn't.
Many points sit on the edges of the square. Making confident predictions from only two numbers means pushing the state to extremes, so the model leans on tanh's limits of −1 and 1. That's a real habit of small RNNs, and it has a cost we'll meet in the section on vanishing gradients.
The axes don't have meanings. It's tempting to read the horizontal axis as "noun-ness" or similar. Don't. Training could just as well have swapped or flipped the axes, and in real models with hundreds of dimensions, information is spread across combinations of units rather than assigned one feature per unit. What matters is where points land relative to each other and to the prediction regions.

Real RNNs have hidden states with hundreds of dimensions, so you can't plot them directly. The picture is the same, though: one point, moving one word at a time, in a space where nearby positions tend to produce similar next-word probabilities.

The same word lands in different places

If the hidden state were just a function of the current word, "the" would always land in the same place. It doesn't.

Two trajectories showing the word the landing in three different places — "the" at the start of a sentence, after "sat on", and after "ran to" lands in three different places, because the old state is part of every update.

Context	Where "the" lands	Top predictions
start of sentence	(0.09, 1.00)	cat, dog, bird
the cat sat on the	(0.97, 0.65)	mat 0.49, rug 0.49
the dog ran to the	(0.96, −0.15)	door 0.48, park 0.48

That last pair is the interesting one. To know whether a mat/rug or a door/park comes next, the model has to remember whether it saw on or to two words ago. That memory is stored entirely in where the point sits. Also notice that "the cat" and "the dog" land on exactly the same point; we'll see why in the next section.

You can drive the model yourself. The explorer below runs the actual trained weights in your browser. Pick a sentence or add words one at a time and watch the point move and the predictions change.

Explore the trained RNN

What Is Token-Specific, and What Is Shared

An RNN has token-specific parameters and shared transformations, and the difference explains a lot about why it works.

Each word owns its embedding. The embedding table has one row per word in the vocabulary. When the model reads "cat", it looks up the cat row, and only that row. Here are some of our model's learned embeddings:

Word	Learned embedding
cat	(−1.30, 0.67)
dog	(−1.30, 0.67)
bird	(−1.48, 0.66)
sat	(−1.31, 0.11)
slept	(−1.29, 0.07)
mat	(0.52, −2.35)
rug	(0.47, −2.37)
the	(1.58, −0.79)

Cat and dog ended up identical to two decimals, bird sits close by, and so do sat/slept and mat/rug. Nobody told the model that these words are similar. They appear in the same company in the training data, so training gave them similar vectors. That's exactly the word2vec idea from The Company Words Keep, arising here as a side effect of predicting the next word. It's also why "the cat" and "the dog" land on the same point.

The recurrent update is shared across positions. $W_x$, $W_h$, and $b$ are reused for every input word at every step. The output layer is also applied at every step, but $W_o$ has a separate row and $c$ a separate bias for each possible next word. So each token has its own input embedding and output-classifier parameters, while the same recurrent rules update the context throughout the sequence.

Sharing buys three things:

Any length works. The same rules apply at step 3 and at step 300.
Lessons transfer. What the model learns about "sat" at position 3 applies wherever "sat" appears, and what it learns about moving from a noun to a verb applies to every noun.
Size doesn't grow with the sentence. Parameters depend on vocabulary size and layer widths, never on how much text you feed in.

Counting the parameters

Let's count for a realistic model: a 10,000-word vocabulary, 300-dimensional embeddings, and a 512-dimensional hidden state.

Piece	Shape	Parameters
Word embeddings	10,000 × 300	3,000,000
$W_x$ (input to hidden)	512 × 300	153,600
$W_h$ (hidden to hidden)	512 × 512	262,144
$b$ (hidden bias)	512	512
$W_o$ (hidden to output)	10,000 × 512	5,120,000
$c$ (output bias)	10,000	10,000
Total		8,546,256

About 8.55 million parameters. Two things stand out:

95% of them live in the two vocabulary-sized tables, the embeddings and $W_o$. The recurrent core that does the "remembering" ($W_x$, $W_h$, and $b$) is only 416,256 parameters, under 5% of the total.
$W_h$ grows with the square of the hidden size. Double the hidden state from 512 to 1,024 and $W_h$ goes from 262,144 to 1,048,576 parameters, four times as many. That quadratic growth is one reason hidden states can't simply be made enormous.

Our toy model has $14 \times 2 = 28$ embedding weights, $2 \times 2 + 2 \times 2 + 2 = 10$ in the recurrent core, and $14 \times 2 + 14 = 42$ in the output layer: 80 parameters, which matches what the code prints.

Three Spaces and the Bridges Between Them

Those 8.5 million weights don't create one giant space. An RNN works with three separate spaces, and the weight matrices are the bridges between them:

Embedding space (300 dimensions in the big example, 2 in ours). Each word is a fixed point here.
Hidden space (512 dimensions, or 2 in ours). This is where the moving point lives.
Output space (one dimension per vocabulary word: 10,000, or 14 in ours). Each point here is a set of scores, one per word.

$W_x$ is the bridge from embedding space into hidden space. $W_h$ is a bridge from hidden space back into itself, which is what makes the network recurrent. $W_o$ is the bridge from hidden space out to the vocabulary.

Word embeddings, hidden-state trajectory, and output probabilities for the trained model — The three spaces of the trained model. Left: the learned word embeddings. Right: the hidden state reading "the cat sat on the". Bottom: the probabilities read off its final position.

A common worry: going from 512 numbers to 10,000 sounds like spreading information thin, as if most of the 10,000 outputs must be empty. It isn't. Each word $w$ gets its own row $u_w$ of $W_o$, a direction in hidden space, and its score is

$$z_w = u_w \cdot h + c_w$$

So all 10,000 scores are computed from the same 512 numbers, each one asking "how far does the hidden state point in my direction?" Nothing is sparse. The scores just aren't independent: they're 10,000 different readings of one 512-dimensional point. Going from 300 to 512 works the same way in the other direction: every hidden unit is a weighted combination of all 300 embedding numbers, a new "view" of the word rather than the same numbers padded with zeros.

From a point to a probability

Let's turn a real hidden state into a prediction. After reading "the", the state is $h = (0.09,\ 1.00)$. The row of $W_o$ for cat is $(2.75,\ 9.08)$ and its bias is 1.37, so its score (called a logit) is

$$z_{\text{cat}} = 2.75(0.09) + 9.08(1.00) + 1.37 = 10.70$$

Doing the same for every word gives dog 10.69, bird 10.29, "." 6.66, and lower scores for everything else. Scores can be any size, so softmax turns them into probabilities that are positive and sum to 1:

$$p_i = \frac{e^{z_i}}{\sum_j e^{z_j}}$$

Subtracting the largest score from every score doesn't change the result and keeps the numbers small. Relative to cat: $e^{0} = 1$, $e^{-0.01} = 0.99$ for dog, $e^{-0.41} = 0.67$ for bird, and the other eleven words add up to only 0.04. The total is 2.70, so

$$p_{\text{cat}} = \frac{1}{2.70} = 0.37, \qquad p_{\text{dog}} = \frac{0.99}{2.70} = 0.37, \qquad p_{\text{bird}} = \frac{0.67}{2.70} = 0.25$$

Bar charts of logits and softmax probabilities after the word the — Logits and softmax probabilities from the real hidden state after "the". Softmax turns gaps between scores into ratios between probabilities.

Here's a nice check. In the training data, sentences start with "the cat" 6 times out of 16, "the dog" 6 times, and "the bird" 4 times: 0.375, 0.375, and 0.25. The model learned roughly those odds.

Why these two functions? tanh for the hidden state because the state is fed back into itself at every step. Bounding it between −1 and 1 keeps it from growing without limit over a long sentence, and being centered on zero lets each unit push in either direction. Softmax for the output because we need a probability distribution over the vocabulary to compare against the word that actually came next.

Training: Learning From Next-Word Prediction

Where do the weights come from? From a game the model plays with itself: read some text, guess the next word, check, adjust.

Every position in every sentence is a training example. "the cat sat on the mat ." gives six of them: the → cat, the cat → sat, the cat sat → on, and so on. Our 16 sentences give 96 predictions per pass.

The loss is cross-entropy: the negative log of the probability the model gave to the word that actually came next.

$$L = -\log p(\text{correct word})$$

If the model gave the right word 0.99, the loss is 0.01. If it gave 0.37, the loss is 0.99. If it gave 0.01, the loss is 4.61. Confident and wrong is punished hard. Training minimizes the average over all 96 predictions.

Our model starts at a loss of 2.770, about what you'd get from guessing uniformly among 14 words ($\ln 14 = 2.64$). After 1,500 passes it reaches 0.484. It can't reach zero, because the data itself is uncertain: after "the", any of three animals can follow. The best possible loss on this corpus is 0.462, so the tiny model is close to perfect.

Backpropagation through time

To improve the weights we need gradients, and the unrolled picture tells us how to get them. The loss at step 4 depends on $h_4$. But $h_4$ was computed from $h_3$, which came from $h_2$, and so on back to the first word. So the gradient flows backward through every earlier step. This is called backpropagation through time (BPTT). It's ordinary backpropagation applied to the unrolled network.

Diagram of gradients flowing backward through an unrolled RNN — Each step's loss sends gradient back down into its hidden state, and from there backward through all earlier steps.

Because the weights are shared, every step contributes to the gradient of the same $W_h$, $W_x$, and $W_o$, and those contributions are added together. Then one gradient-descent step nudges each weight downhill, exactly as in Introduction to Machine Learning. In PyTorch, all of this happens inside one call to loss.backward().

Generalization by Proximity

Remember the held-out sentence. The model never saw "the bird sat". What does it do with it?

Trajectories for the cat sat and the bird sat ending in the same region — "the bird" lands next to "the cat", and "sat" carries both into the region that predicts "on". The model has never seen "the bird sat" during training.

Prefix	Hidden state	Top predictions
the cat	(−0.64, 0.16)	ran 0.35, slept 0.35, sat 0.28
the bird	(−0.72, −0.13)	slept 0.43, ran 0.43, sat 0.10
the cat sat	(0.04, −0.98)	on 0.99
the dog sat	(0.04, −0.98)	on 0.99
the bird sat	(0.14, −1.00)	on 0.98

Two things happened, and both are reasonable.

First, the model noticed the gap. It gives sat only 0.10 after "the bird", versus 0.28 after "the cat". In its training data, birds never sat.

Second, when it does see "the bird sat", it reuses what it knows. Bird shares company with cat and dog (they all slept and ran in the same places), so its embedding is close to theirs: 0.18 away from cat, compared with 3.53 between cat and mat. A nearby embedding produces a nearby hidden state, "the bird" lands near "the cat", and the shared update for sat carries it into the same neighborhood. "the bird sat" ends up 0.10 away from "the cat sat", inside the region that predicts on, and the model predicts on with probability 0.98.

This is how RNNs, and neural language models in general, handle combinations they've never seen. Similar contexts land in similar places, and nearby places make similar predictions. It cuts both ways, though. Proximity transfers patterns whether or not they're true, which is one reason language models can be fluent and confidently wrong.

Where Simple RNNs Struggle

Vanishing and exploding gradients

Backpropagation through time multiplies. To see how much the final state depends on a state $k$ steps earlier, you chain one factor per step:

$$\frac{\partial h_T}{\partial h_{T-k}} = \frac{\partial h_T}{\partial h_{T-1}} \cdot \frac{\partial h_{T-1}}{\partial h_{T-2}} \cdots \frac{\partial h_{T-k+1}}{\partial h_{T-k}}$$

and each factor is the same kind of thing:

$$\frac{\partial h_t}{\partial h_{t-1}} = \text{diag}\left(1 - h_t^2\right) W_h$$

The first part is the slope of tanh at each unit. It's at most 1, and it's close to 0 when a unit is saturated near −1 or 1. The second part is the recurrent matrix, the same one at every step. Multiply many of these together and the result behaves like a number raised to a power: if the factors are typically smaller than 1, the product shrinks exponentially; if they're typically larger, it grows exponentially.

Let's measure it on our trained model, using a 28-word stream of four sentences back to back:

grad_check.py

# Run after the script above. How much does the last hidden state still
# depend on the state k steps earlier? Each step back multiplies by the
# Jacobian dh_t/dh_{t-1} = diag(1 - h_t^2) @ W_h.
stream = ("the cat sat on the mat . the dog ran to the park . "
          "the bird slept on the rug . the cat ran to the door .")
with torch.no_grad():
    _, states = model(torch.tensor([[stoi[w] for w in stream.split()]]))
h = states[0]                         # (28, 2): one point per word
W_h = model.W_h.weight.detach()

print(f"largest stretch W_h can apply: {torch.linalg.matrix_norm(W_h, ord=2):.2f}")
print(f"average tanh slope along the stream: {(1 - h ** 2).mean():.2f}")
J = torch.eye(2)
for k in range(1, 9):
    J = J @ (torch.diag(1 - h[-k] ** 2) @ W_h)
    print(f"{k} step(s) back: {torch.linalg.matrix_norm(J, ord=2):.1e}")

output

largest stretch W_h can apply: 3.39
average tanh slope along the stream: 0.27
1 step(s) back: 6.1e-01
2 step(s) back: 9.1e-04
3 step(s) back: 3.0e-03
4 step(s) back: 1.4e-06
5 step(s) back: 2.6e-06
6 step(s) back: 1.5e-06
7 step(s) back: 2.7e-07
8 step(s) back: 4.2e-08

$W_h$ on its own could stretch a vector by up to 3.39, which sounds like a recipe for explosion. But the average tanh slope along the way is only 0.27, because so many states sit at the edges of the square. Four words back, the signal has shrunk to about one millionth. Saturation doesn't focus the gradient; it starves it.

The same measurement on a 64-unit version of the model, trained on the same sentences, shows the more typical pattern:

Log-scale plot of gradient norm versus steps back in time for two trained models — How strongly the final hidden state depends on each earlier one, measured on two trained models (log scale). The 64-unit model first grows (2.6 at one step back) and then fades roughly exponentially; the 2-unit model collapses within a few steps and drops off the bottom of the chart after 11.

The bigger model's signal first grows, to 2.6 one step back, then fades steadily, falling below $10^{-4}$ by 18 steps back. That's the vanishing gradient problem: the training signal from a prediction can't reach words far in the past, so the model can't learn that they matter. Consider "The keys that I left on the kitchen table yesterday ___". Choosing are over is depends on keys, nine words earlier. If almost no gradient flows that far back, the model never learns the connection.

The opposite failure, exploding gradients, happens when the factors are mostly larger than 1. Gradients become huge, a single update throws the weights far off, and the loss jumps or turns into NaN. The standard fix is gradient clipping: if the gradient's overall length exceeds a threshold, scale it down to that length before the update. Our training loop does this with clip_grad_norm_(model.parameters(), 1.0). Clipping tames explosions, but it does nothing for vanishing gradients; you can't rescale a signal that's already gone.

Everything squeezed into one point

There's a second, simpler limit. The hidden state is a fixed-size vector, and it's overwritten at every step. Everything the model knows about the past has to fit in that one point, and each new word partly erases what was there. A basic RNN also has no way to look back at a specific earlier word and decide it matters more than the others. Whatever survived in the hidden state is all it has.

What Came Next

Each of the main successors to the basic RNN attacks one of these limits.

LSTM (long short-term memory, Hochreiter and Schmidhuber, 1997) adds a second vector, the cell state, plus three gates: small sigmoid layers that output values between 0 and 1 and decide what to forget, what to write, and what to reveal. The cell state is updated mostly by adding and scaling rather than by squashing through tanh with a full matrix multiply at every step. When the forget gate stays near 1, the gradient can travel back along the cell state for many steps without vanishing.

GRU (gated recurrent unit, Cho et al., 2014) simplifies the LSTM to two gates and a single state vector. It often performs about as well with fewer parameters.

Attention (Bahdanau, Cho, and Bengio, 2014) removes the one-point bottleneck. Instead of predicting from the final hidden state alone, the model keeps every hidden state and, at each step, computes learned weights over them: an explicit answer to "which earlier words matter right now?" This is the importance mechanism a basic RNN lacks. Transformers (Vaswani et al., 2017) went further, dropping recurrence entirely and building everything from attention, which also lets them process all positions in parallel.

The moving-point picture still applies to LSTMs and GRUs. Their state is still a point moving through space one word at a time; the gates just give the network much finer control over how far, and in which directions, each word is allowed to move it.

The Full Implementation

Here's the complete program behind every number and figure in this tutorial. It writes the RNN loop out by hand so you can see the equation in the code. PyTorch's built-in nn.RNN runs the same loop in optimized form; it keeps two bias vectors (bias_ih_l0 and bias_hh_l0) that simply add together into our single $b$.

tiny_rnn.py

import torch
import torch.nn as nn
import torch.nn.functional as F

torch.manual_seed(2)

# ---------------------------------------------------------------
# 1. A toy corpus: every "the    the  ."
#    except that "the bird sat ..." never appears. We hold it out
#    to test generalization later.
# ---------------------------------------------------------------
subjects = ["cat", "dog", "bird"]
verbs = {"sat": "on", "slept": "on", "ran": "to"}       # verb -> preposition
places = {"on": ["mat", "rug"], "to": ["door", "park"]}  # preposition -> places

sentences = [
    f"the {s} {v} {p} the {o} ."
    for s in subjects
    for v, p in verbs.items()
    for o in places[p]
    if (s, v) != ("bird", "sat")
]
vocab = sorted({w for s in sentences for w in s.split()})
stoi = {w: i for i, w in enumerate(vocab)}
data = torch.tensor([[stoi[w] for w in s.split()] for s in sentences])
inputs, targets = data[:, :-1], data[:, 1:]  # predict word t+1 from words 1..t

print(f"{len(sentences)} sentences, vocabulary of {len(vocab)}: {vocab}")


# ---------------------------------------------------------------
# 2. The model: h_t = tanh(W_h h_{t-1} + W_x x_t + b),  logits = W_o h_t + c
# ---------------------------------------------------------------
class TinyRNN(nn.Module):
    def __init__(self, vocab_size, emb_dim=2, hidden_dim=2):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.embed = nn.Embedding(vocab_size, emb_dim)            # one vector per word
        self.W_x = nn.Linear(emb_dim, hidden_dim)                 # shared; its bias is b
        self.W_h = nn.Linear(hidden_dim, hidden_dim, bias=False)  # shared
        self.W_o = nn.Linear(hidden_dim, vocab_size)              # shared; its bias is c

    def forward(self, tokens):
        batch, steps = tokens.shape
        h = torch.zeros(batch, self.hidden_dim)  # h_0: start at the origin
        states = []
        for t in range(steps):                   # same weights at every step
            x_t = self.embed(tokens[:, t])
            h = torch.tanh(self.W_h(h) + self.W_x(x_t))
            states.append(h)
        states = torch.stack(states, dim=1)      # (batch, steps, hidden_dim)
        return self.W_o(states), states          # logits for every step, plus the path


model = TinyRNN(len(vocab))
n_params = sum(p.numel() for p in model.parameters())
print(f"parameters: {n_params}")

# ---------------------------------------------------------------
# 3. Train: next-word prediction with cross-entropy, backprop through time
# ---------------------------------------------------------------
optimizer = torch.optim.AdamW(model.parameters(), lr=0.05, weight_decay=0.01)

for epoch in range(1, 1501):
    logits, _ = model(inputs)
    loss = F.cross_entropy(logits.reshape(-1, len(vocab)), targets.reshape(-1))
    optimizer.zero_grad()
    loss.backward()                                            # BPTT happens here
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)    # gradient clipping
    optimizer.step()
    if epoch in (1, 100, 500, 1500):
        print(f"epoch {epoch:4d}  loss {loss.item():.3f}")


# ---------------------------------------------------------------
# 4. Look inside: where does each prefix land, and what comes next?
# ---------------------------------------------------------------
@torch.no_grad()
def predict(prompt, k=3):
    tokens = torch.tensor([[stoi[w] for w in prompt.split()]])
    logits, states = model(tokens)
    probs = logits[0, -1].softmax(-1)
    top = probs.topk(k)
    h = states[0, -1].tolist()
    guesses = "  ".join(f"{vocab[i]} {p:.2f}" for p, i in zip(top.values.tolist(), top.indices.tolist()))
    print(f"{prompt:<22} h = ({h[0]:+.2f}, {h[1]:+.2f})   next: {guesses}")


print()
for prompt in ["the", "the cat", "the cat sat", "the cat sat on", "the cat sat on the",
               "the dog ran to the", "the cat sat on the mat"]:
    predict(prompt)

print("\nGeneralization (the model never saw 'the bird sat'):")
for prompt in ["the cat", "the bird", "the cat sat", "the dog sat", "the bird sat"]:
    predict(prompt)

Running it prints: