Skip to content

deepbox/nn

Recurrent & Packed Sequences

RNN, LSTM and GRU layers, and helpers that pack and unpack variable-length sequences.

  • Sequences

Overview

  • The input layout is [batch, seq, feature] by default (batchFirst: true). PyTorch defaults to sequence first. Other options are numLayers, bidirectional, bias, and nonlinearity for RNN.
  • forward(x) returns the outputs of every step. forwardWithState(x, hx?) also returns the final hidden state, as [output, hN], and for LSTM as [output, [hN, cN]].
  • Layers compute in their parameter dtype, so integer input is accepted. The recurrent layers run host kernels, so their weights stay on the CPU after model.to(device).
  • packSequence stores sequences of different lengths in one flat tensor sorted by length, and unpackSequence and padPackedSequence reverse it. Packed data is plain tensors outside autograd, and the recurrent layers take padded tensors, not PackedSequence objects.

Types

2

PackedSequence

type

A packed representation of variable-length sequences.

Stores sequence data sorted by length (longest first) in a flat tensor, along with batch sizes at each timestep.

  • readonly data: Tensor

    Packed data tensor of shape (totalElements, features)

  • readonly batchSizes: number[]

    Number of active sequences at each timestep

  • readonly sortedIndices: number[]

    Original index of each sequence, listed in sorted (longest first) order

  • readonly unsortedIndices: number[]

    Position of each original sequence within the sorted order (inverse of sortedIndices)

  • readonly featureSize: number

    Feature dimension size

  • readonly squeezeFeature?: boolean

    True when the packed sequences were 1D, so unpackSequence returns 1D tensors of shape (seqLen) instead of (seqLen, 1).

RNNNonlinearity

type
type RNNNonlinearity = "tanh" | "relu";

Classes

3

GRU

extends Module

GRU (Gated Recurrent Unit) layer.

Gate weights are stacked in PyTorch's order (reset, update, new) and the reset gate is applied to the hidden projection only, as in PyTorch. Input and output layouts, parameter names, initialization and dtype handling follow RNN.

Example
import { GRU } from 'deepbox/nn'; const gru = new GRU(4, 8);const output = gru.forward(x);
Members (13)
  • constructor(inputSize: number, hiddenSize: number, options?: { readonly numLayers?: number; readonly bias?: boolean; readonly batchFirst?: boolean; readonly bidirectional?: boolean; readonly dtype?: "float32" | "float64"; })
  • readonly inputSize: number

    Number of features of each input step.

  • readonly hiddenSize: number

    Number of features of the hidden state.

  • readonly numLayers: number

    Number of stacked layers.

  • readonly batchFirst: boolean

    Whether input and output are (batch, seq, feature) rather than (seq, batch, feature).

  • readonly bidirectional: boolean

    Whether each layer also runs over the sequence backwards.

  • forward(input: GradTensor, hx?: AnyTensor): GradTensor
  • forward(input: Tensor, hx?: AnyTensor): AnyTensor
  • forward(...inputs: AnyTensor[]): AnyTensor
  • forwardWithState(input: GradTensor, hx?: AnyTensor): [ GradTensor, GradTensor ]

    Run the layer and also return the final hidden state.

  • forwardWithState(input: Tensor, hx?: AnyTensor): [ AnyTensor, AnyTensor ]
  • forwardWithState(input: AnyTensor, hx?: AnyTensor): [ AnyTensor, AnyTensor ]
  • toString(): string

LSTM

extends Module

LSTM (Long Short-Term Memory) layer.

Gate weights are stacked in PyTorch's order (input, forget, cell, output). Input and output layouts, parameter names, initialization and dtype handling follow RNN. forwardWithState returns the output and the final hidden and cell states [output, [hN, cN]].

Example
import { LSTM } from 'deepbox/nn'; const lstm = new LSTM(4, 8);const [output, [hN, cN]] = lstm.forwardWithState(x);
Members (13)
  • constructor(inputSize: number, hiddenSize: number, options?: { readonly numLayers?: number; readonly bias?: boolean; readonly batchFirst?: boolean; readonly bidirectional?: boolean; readonly dtype?: "float32" | "float64"; })
  • readonly inputSize: number

    Number of features of each input step.

  • readonly hiddenSize: number

    Number of features of the hidden state.

  • readonly numLayers: number

    Number of stacked layers.

  • readonly batchFirst: boolean

    Whether input and output are (batch, seq, feature) rather than (seq, batch, feature).

  • readonly bidirectional: boolean

    Whether each layer also runs over the sequence backwards.

  • forward(input: GradTensor, hx?: AnyTensor, cx?: AnyTensor): GradTensor
  • forward(input: Tensor, hx?: AnyTensor, cx?: AnyTensor): AnyTensor
  • forward(...inputs: AnyTensor[]): AnyTensor
  • forwardWithState(input: GradTensor, hx?: AnyTensor, cx?: AnyTensor): [ GradTensor, [ GradTensor, GradTensor ] ]

    Run the layer and also return the final hidden and cell states.

  • forwardWithState(input: Tensor, hx?: AnyTensor, cx?: AnyTensor): [ AnyTensor, [ AnyTensor, AnyTensor ] ]
  • forwardWithState(input: AnyTensor, hx?: AnyTensor, cx?: AnyTensor): [ AnyTensor, [ AnyTensor, AnyTensor ] ]
  • toString(): string

RNN

extends Module

Elman RNN layer with tanh or relu nonlinearity.

Computes h_t = act(W_ih x_t + b_ih + W_hh h_(t-1) + b_hh) for every layer and direction. Input is (batch, seq, feature) (or (seq, batch, feature) with batchFirst: false), or unbatched (seq, feature). As in PyTorch, all weights and biases are drawn from the uniform distribution U(-1/sqrt(hiddenSize), 1/sqrt(hiddenSize)). Parameters are named like PyTorch's (weight_ih_l0, weight_hh_l0, bias_ih_l0, bias_hh_l0, and a _reverse suffix for the backward direction). The layer computes in the parameter dtype and casts the input (and the initial state) to it.

forward returns the output sequence (batch, seq, hidden * directions); forwardWithState also returns the final hidden state (layers * directions, batch, hidden). A GradTensor input gives GradTensor results. A plain Tensor input gives GradTensor results that track the weights while they require grad and gradient tracking is on, and plain tensors otherwise (inside noGrad() or with frozen weights).

Example
import { RNN } from 'deepbox/nn';import { randn } from 'deepbox/ndarray'; const rnn = new RNN(4, 8, { numLayers: 2 });const x = randn([2, 5, 4]); // (batch, seq, feature)const output = rnn.forward(x);
Members (14)
  • constructor(inputSize: number, hiddenSize: number, options?: { readonly numLayers?: number; readonly nonlinearity?: RNNNonlinearity; readonly bias?: boolean; readonly batchFirst?: boolean; readonly bidirectional?: boolean; readonly dtype?: "float32" | "float64"; })
  • readonly inputSize: number

    Number of features of each input step.

  • readonly hiddenSize: number

    Number of features of the hidden state.

  • readonly numLayers: number

    Number of stacked layers.

  • readonly nonlinearity: RNNNonlinearity

    Activation applied at every step.

  • readonly batchFirst: boolean

    Whether input and output are (batch, seq, feature) rather than (seq, batch, feature).

  • readonly bidirectional: boolean

    Whether each layer also runs over the sequence backwards.

  • forward(input: GradTensor, hx?: AnyTensor): GradTensor
  • forward(input: Tensor, hx?: AnyTensor): AnyTensor
  • forward(...inputs: AnyTensor[]): AnyTensor
  • forwardWithState(input: GradTensor, hx?: AnyTensor): [ GradTensor, GradTensor ]

    Run the layer and also return the final hidden state.

  • forwardWithState(input: Tensor, hx?: AnyTensor): [ AnyTensor, AnyTensor ]
  • forwardWithState(input: AnyTensor, hx?: AnyTensor): [ AnyTensor, AnyTensor ]
  • toString(): string

Functions

4

packPaddedSequence

function packPaddedSequence(input: Tensor, lengths: readonly number[], enforcesSorted?: boolean): PackedSequence;

Pack a padded 3D tensor into a PackedSequence.

Takes a padded tensor of shape (batch, maxLen, features) and corresponding lengths, and creates a packed representation. The input dtype is kept.

Parameters

input
Tensor
Padded tensor of shape (batch, maxLen, features)
lengths
readonly number[]
Actual lengths of each sequence in the batch
enforcesSortedoptional
boolean
If true, the lengths must already be sorted descending (an error is thrown otherwise)

Returns

PackedSequencePacked sequence representation

Throws

  • ShapeErrorIf input is not 3D
  • InvalidParameterErrorIf lengths has the wrong number of entries or an entry is not an integer in [1, maxLen]
Example
import { packPaddedSequence } from 'deepbox/nn';import { tensor } from 'deepbox/ndarray'; const padded = tensor([[[1, 2], [3, 4]], [[5, 6], [0, 0]]]);const packed = packPaddedSequence(padded, [2, 1]);

packSequence

function packSequence(sequences: readonly Tensor[], enforcesSorted?: boolean): PackedSequence;

Pack a list of variable-length tensors into a PackedSequence.

Sequences are sorted by length (longest first) and interleaved so that at each timestep only active sequences are present. All sequences are converted to the dtype of the first one.

Parameters

sequences
readonly Tensor[]
Array of 1D or 2D tensors. If 2D, shape is (seqLen, features). If 1D, treated as (seqLen, 1).
enforcesSortedoptional
boolean
If true, the sequences must already be sorted by descending length (an error is thrown otherwise) and no sorting is done. Default: false.

Returns

PackedSequencePacked sequence representation

Throws

  • InvalidParameterErrorIf sequences is empty, a sequence has zero length, or enforcesSorted is set but the lengths are not descending
  • ShapeErrorIf a sequence is not 1D or 2D, or feature sizes differ
  • DTypeErrorIf a sequence is a string tensor
Example
import { packSequence } from 'deepbox/nn';import { tensor } from 'deepbox/ndarray'; const seqs = [  tensor([[1, 2], [3, 4], [5, 6]]),  tensor([[7, 8]]),];const packed = packSequence(seqs);

padPackedSequence

function padPackedSequence(packed: PackedSequence, totalLength?: number, paddingValue?: number): [ Tensor, number[] ];

Pad a PackedSequence to a dense tensor.

Creates a padded 3D tensor of shape (batchSize, maxLen, features) filled with paddingValue after the end of shorter sequences. The returned tensor has sequences in their original order.

Parameters

packed
PackedSequence
Packed sequence to pad
totalLengthoptional
number
Optional total length to pad to (defaults to max sequence length). Must be at least the length of the longest sequence.
paddingValueoptional
number
Value written after the end of each sequence (default: 0)

Returns

[ Tensor, number[] ]Tuple of [paddedTensor, lengths]

Throws

  • InvalidParameterErrorIf totalLength is not an integer or is shorter than the longest sequence, or the packed sequence is inconsistent
Example
import { packSequence, padPackedSequence } from 'deepbox/nn';import { tensor } from 'deepbox/ndarray'; const seqs = [tensor([[1, 2], [3, 4]]), tensor([[5, 6]])];const packed = packSequence(seqs);const [padded, lengths] = padPackedSequence(packed);// padded.shape = [2, 2, 2], lengths = [2, 1]

unpackSequence

function unpackSequence(packed: PackedSequence): [ Tensor[], number[] ];

Unpack a PackedSequence back to a list of tensors.

Returns the sequences in the original order (before sorting). Each tensor has shape (seqLen, features), or (seqLen) when the packed sequences were 1D.

Parameters

packed
PackedSequence
Packed sequence to unpack

Returns

[ Tensor[], number[] ]Tuple of [sequences, lengths] where sequences are in original order

Throws

  • InvalidParameterErrorIf the packed sequence is inconsistent
  • ShapeErrorIf the packed data size does not match batchSizes
Example
import { packSequence, unpackSequence } from 'deepbox/nn';import { tensor } from 'deepbox/ndarray'; const seqs = [tensor([[1, 2], [3, 4]]), tensor([[5, 6]])];const packed = packSequence(seqs);const [unpacked, lengths] = unpackSequence(packed);
nn-recurrent.ts
import { randn, tensor } from "deepbox/ndarray";import { GRU, LSTM, RNN, packSequence, unpackSequence } from "deepbox/nn"; const x = randn([2, 5, 3]); // [batch, seq, feature] const lstm = new LSTM(3, 4);const [output, [hN, cN]] = lstm.forwardWithState(x);console.log(output.shape, hN.shape, cN.shape); // [2, 5, 4] [1, 2, 4] [1, 2, 4] console.log(new RNN(3, 4, { bidirectional: true }).forward(x).shape); // [2, 5, 8]console.log(new GRU(3, 4, { numLayers: 2 }).forward(x).shape); // [2, 5, 4] const packed = packSequence([tensor([[1], [2], [3]]), tensor([[4], [5]])]);console.log(packed.batchSizes);console.log(unpackSequence(packed)[1]); // original lengths

Found a mistake? Open an issue.