Implementation:Tensorflow Tfjs Recurrent Layers

Knowledge Sources	Tensorflow_Tfjs TensorFlow.js API
Domains	Deep_Learning, Layers_API
Last Updated	2026-02-10 06:00 GMT

Overview

The recurrent.ts module implements all recurrent neural network (RNN) layers for the TensorFlow.js Layers API. It provides the core rnn() step function, the abstract RNN and RNNCell base classes, and concrete implementations: SimpleRNN/SimpleRNNCell, GRU/GRUCell, LSTM/LSTMCell, and StackedRNNCells.

The architecture follows a cell/layer separation pattern: each cell class implements a single timestep computation, while the corresponding layer class uses the rnn() function to iterate the cell over the time dimension of the input. This allows cells to be used independently or stacked via StackedRNNCells.

Key features include support for return sequences, return state, bidirectional processing (via the wrapper module), masking, dropout, and recurrent dropout.

Code Reference

Source Location

tfjs-layers/src/layers/recurrent.ts (View on GitHub)

Key Signatures

// Input standardization
export function standardizeArgs(
    inputs: Tensor | Tensor[] | SymbolicTensor | SymbolicTensor[],
    initialState: Tensor | Tensor[] | SymbolicTensor | SymbolicTensor[],
    constants: Tensor | Tensor[] | SymbolicTensor | SymbolicTensor[],
    numConstants?: number
): { inputs, initialState, constants };

// Core RNN step function
export function rnn(
    stepFunction: RnnStepFunction,
    inputs: Tensor,
    initialStates: Tensor[],
    goBackwards?: boolean,
    mask?: Tensor,
    constants?: Tensor[],
    unroll?: boolean,
    needPerStepOutputs?: boolean
): [Tensor, Tensor, Tensor[]];

// Abstract base classes
export abstract class RNNCell extends Layer {
  stateSize: number | number[];
  dropoutMask: Tensor | Tensor[];
  recurrentDropoutMask: Tensor | Tensor[];
}

export class RNN extends Layer {
  static className = 'RNN';
  readonly cell: RNNCell;
  readonly returnSequences: boolean;
  readonly returnState: boolean;
  readonly goBackwards: boolean;
  readonly stateful: boolean;
  states_: Tensor[];
  constructor(args: RNNLayerArgs);
  call(inputs: Tensor | Tensor[], kwargs: Kwargs): Tensor | Tensor[];
  resetStates(states?: Tensor | Tensor[], training?: boolean): void;
}

// Concrete cells and layers
export class SimpleRNNCell extends RNNCell { static className = 'SimpleRNNCell'; }
export class SimpleRNN extends RNN { static className = 'SimpleRNN'; }
export class GRUCell extends RNNCell { static className = 'GRUCell'; }
export class GRU extends RNN { static className = 'GRU'; }
export class LSTMCell extends RNNCell { static className = 'LSTMCell'; }
export class LSTM extends RNN { static className = 'LSTM'; }
export class StackedRNNCells extends RNNCell { static className = 'StackedRNNCells'; }

// Dropout mask generation
export function generateDropoutMask(args: {
  ones: () => Tensor, rate: number, training?: boolean, count?: number
}): Tensor | Tensor[];

Import

import { RNN, RNNCell, SimpleRNN, GRU, LSTM, rnn, StackedRNNCells } from './layers/recurrent';

Cell Details

Cell	State Size	Key Config	Gates
SimpleRNNCell	units	units, activation (tanh), useBias, dropout, recurrentDropout	Single gate: h = activation(x*W + h_prev*U + b)
GRUCell	units	units, activation (tanh), recurrentActivation (hardSigmoid), resetAfter, implementation	Update (z), Reset (r), New (h~)
LSTMCell	[units, units]	units, activation (tanh), recurrentActivation (hardSigmoid), unitForgetBias, implementation	Input (i), Forget (f), Cell (c), Output (o)
StackedRNNCells	per-cell sizes	cells (array of RNNCells)	Delegates to each cell in sequence

I/O Contract

rnn() Function

• Input: A step function, rank-3+ input tensor [batch, time, ...], initial states, and optional mask/constants.
• Output: [lastOutput, outputs, newStates] where outputs is populated only if needPerStepOutputs=true.
• Behavior: Transposes input to time-major, iterates over timesteps calling the step function, optionally applies mask, supports reverse iteration (goBackwards).

RNN Layer

• Input: Rank-3 tensor [batch, timesteps, features]. Optional initialState and mask in kwargs.
• Output:
  • returnSequences=false: [batch, units]
  • returnSequences=true: [batch, timesteps, units]
  • returnState=true: Additional state tensors appended (for LSTM: [h, c])

Stateful Mode

When stateful=true, the RNN maintains states across batches. States are reset explicitly via resetStates(). Requires fixed batch size via batchInputShape.

Usage Examples

import * as tf from '@tensorflow/tfjs';

// Simple LSTM model for sequence classification
const model = tf.sequential();
model.add(tf.layers.lstm({
  inputShape: [100, 32],  // 100 timesteps, 32 features
  units: 64,
  returnSequences: false
}));
model.add(tf.layers.dense({ units: 10, activation: 'softmax' }));

// Stacked GRU with return sequences
const stackedModel = tf.sequential();
stackedModel.add(tf.layers.gru({
  inputShape: [50, 16],
  units: 128,
  returnSequences: true,
  dropout: 0.2,
  recurrentDropout: 0.2
}));
stackedModel.add(tf.layers.gru({ units: 64, returnSequences: false }));

// Using StackedRNNCells with the RNN layer
const cells = [
  tf.layers.lstmCell({ units: 128 }),
  tf.layers.lstmCell({ units: 64 })
];
const rnnLayer = tf.layers.rnn({
  cell: tf.layers.stackedRNNCells({ cells }),
  returnSequences: true
});

// LSTM returning state for encoder-decoder
const encoder = tf.layers.lstm({
  units: 256,
  returnState: true  // returns [output, h, c]
});

Related Pages

• Tensorflow_Tfjs_ConvLSTM2D_Layer - Convolutional LSTM that extends the RNN/LSTMCell architecture
• Tensorflow_Tfjs_Wrapper_Layers - Bidirectional wrapper that runs an RNN in both directions
• Tensorflow_Tfjs_Backend_Operations - K.dropout and K.sliceAlongFirstAxis used in RNN computation
• Tensorflow_Tfjs_Exports_Layers - Public API factory functions (tf.layers.lstm, tf.layers.gru, etc.)
• Tensorflow_Tfjs_Layer_Variables - batchGetValue/batchSetValue used for state management in StackedRNNCells