Skip to content
28intermediate

Recurrent Neural Network Layers

Runs the RNN, LSTM and GRU layers on a small batch of sequences and prints the output shapes. It also shows stacked layers, unbatched input and the weight count of each layer type.

What you will learn

  • Use deepbox/ndarray for tensor, noGrad.
  • Use deepbox/nn for RNN, LSTM, GRU.
  • Runs the RNN, LSTM and GRU layers on a small batch of sequences and prints the output shapes. It also shows stacked layers, unbatched input and the weight count of each layer type.

Source

index.ts
/** * Example 28: Recurrent Neural Network Layers * * RNN, LSTM and GRU layers on a small batch of sequences, with output shapes * and parameter counts. Recurrent layers carry a hidden state across time steps. * * The forward passes are inference only, so they run inside noGrad() and * return plain tensors. Without it, the layers return a GradTensor that * tracks the weights. Both kinds of tensor have .shape. */ import { noGrad, tensor } from "deepbox/ndarray";import { GRU, LSTM, RNN } from "deepbox/nn"; console.log("=== Recurrent Neural Network Layers ===\n"); // ---------------------------------------------------------------------------// Part 1: Simple RNN// ---------------------------------------------------------------------------console.log("--- Part 1: Simple RNN ---"); // RNN(inputSize, hiddenSize, options)// Input shape (batchFirst=true): (batch, seqLen, inputSize)const rnn = new RNN(4, 8, { batchFirst: true });console.log("RNN(inputSize=4, hiddenSize=8, batchFirst=true)"); // Batch of 2 sequences, each with 3 time steps and 4 featuresconst rnnInput = tensor([  [    [1, 2, 3, 4],    [5, 6, 7, 8],    [9, 10, 11, 12],  ],  [    [13, 14, 15, 16],    [17, 18, 19, 20],    [21, 22, 23, 24],  ],]);console.log(`Input shape:  [${rnnInput.shape.join(", ")}]`); const rnnOut = noGrad(() => rnn.forward(rnnInput));console.log(`Output shape: [${rnnOut.shape.join(", ")}]`);console.log("  The output holds the hidden state at every time step\n"); // ---------------------------------------------------------------------------// Part 2: LSTM (Long Short-Term Memory)// ---------------------------------------------------------------------------console.log("--- Part 2: LSTM ---"); // LSTM adds cell state for better long-range dependenciesconst lstm = new LSTM(4, 8, { batchFirst: true });console.log("LSTM(inputSize=4, hiddenSize=8, batchFirst=true)");console.log(`Input shape:  [${rnnInput.shape.join(", ")}]`); const lstmOut = noGrad(() => lstm.forward(rnnInput));console.log(`Output shape: [${lstmOut.shape.join(", ")}]`);console.log("  LSTM uses forget/input/output gates for selective memory\n"); // ---------------------------------------------------------------------------// Part 3: GRU (Gated Recurrent Unit)// ---------------------------------------------------------------------------console.log("--- Part 3: GRU ---"); // GRU is a simplified version of LSTM with fewer parametersconst gru = new GRU(4, 8, { batchFirst: true });console.log("GRU(inputSize=4, hiddenSize=8, batchFirst=true)");console.log(`Input shape:  [${rnnInput.shape.join(", ")}]`); const gruOut = noGrad(() => gru.forward(rnnInput));console.log(`Output shape: [${gruOut.shape.join(", ")}]`);console.log("  GRU uses reset/update gates, fewer parameters than LSTM\n"); // ---------------------------------------------------------------------------// Part 4: Multi-layer RNN// ---------------------------------------------------------------------------console.log("--- Part 4: Multi-Layer Stacking ---"); const deepRnn = new RNN(4, 16, { numLayers: 2, batchFirst: true });console.log("RNN(inputSize=4, hiddenSize=16, numLayers=2)");console.log(`Input shape:  [${rnnInput.shape.join(", ")}]`); const deepOut = noGrad(() => deepRnn.forward(rnnInput));console.log(`Output shape: [${deepOut.shape.join(", ")}]`);console.log("  2-layer RNN feeds the first layer's hidden states into a second layer\n"); // ---------------------------------------------------------------------------// Part 5: Unbatched (single sequence) input// ---------------------------------------------------------------------------console.log("--- Part 5: Unbatched Input ---"); const singleSeq = tensor([  [1, 2, 3, 4],  [5, 6, 7, 8],  [9, 10, 11, 12],]);console.log("Single sequence (no batch dim):");console.log(`Input shape:  [${singleSeq.shape.join(", ")}]`); const singleOut = noGrad(() => rnn.forward(singleSeq));console.log(`Output shape: [${singleOut.shape.join(", ")}]`);console.log("  2D input is treated as unbatched sequence\n"); // ---------------------------------------------------------------------------// Part 6: Parameter counts// ---------------------------------------------------------------------------console.log("--- Part 6: Weight Counts ---");// Count scalar weights, not parameter tensors: every layer has 4 tensors// (input weights, hidden weights and two biases), but the tensors differ in size.const countWeights = (layer: { parameters(): Iterable<{ size: number }> }): number =>  Array.from(layer.parameters()).reduce((total, p) => total + p.size, 0); const rnnWeights = countWeights(rnn);const lstmWeights = countWeights(lstm);const gruWeights = countWeights(gru);console.log(`RNN  weights: ${rnnWeights}  (1 transform per step)`);console.log(`LSTM weights: ${lstmWeights}  (4 gates, ${lstmWeights / rnnWeights}x the RNN)`);console.log(`GRU  weights: ${gruWeights}  (3 gates, ${gruWeights / rnnWeights}x the RNN)`); console.log("\n=== Recurrent Layers Complete ===");

Output

Console output
Console output only: output shapes for each layer type, and a comparison of the number of weights in RNN, LSTM and GRU layers of the same size.
The forward passes run inside `noGrad()`, so they return plain tensors. Outside `noGrad()` the layers return a `GradTensor` that tracks the weights.