Implementation:Tensorflow Tfjs Wrapper Layers

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

Overview

The wrappers.ts module implements wrapper layers that augment the functionality of existing layers. The module provides an abstract Wrapper base class and two concrete implementations:

• TimeDistributed: Applies a layer independently to each timestep in a sequence input. Uses the rnn() step function internally to iterate over the time dimension, applying the wrapped layer at each step.
• Bidirectional: Wraps an RNN layer to process sequences in both forward and backward directions, then merges the results using a configurable merge mode (concat, sum, ave, mul, or null for separate outputs).

The Wrapper base class delegates weight management, trainability, and serialization to the wrapped layer, ensuring transparent behavior in the model graph.

Code Reference

Source Location

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

Key Signatures

export abstract class Wrapper extends Layer {
  readonly layer: Layer;
  constructor(args: WrapperLayerArgs);
  get trainable(): boolean;
  set trainable(value: boolean);
  get trainableWeights(): LayerVariable[];
  get nonTrainableWeights(): LayerVariable[];
  getWeights(): Tensor[];
  setWeights(weights: Tensor[]): void;
  getConfig(): serialization.ConfigDict;
  static fromConfig<T>(cls, config, customObjects?): T;
}

export class TimeDistributed extends Wrapper {
  static className = 'TimeDistributed';
  constructor(args: WrapperLayerArgs);
  build(inputShape: Shape | Shape[]): void;
  computeOutputShape(inputShape: Shape | Shape[]): Shape | Shape[];
  call(inputs: Tensor | Tensor[], kwargs: Kwargs): Tensor | Tensor[];
}

export class Bidirectional extends Wrapper {
  static className = 'Bidirectional';
  readonly mergeMode: BidirectionalMergeMode;
  readonly forwardLayer: RNN;
  readonly backwardLayer: RNN;
  readonly returnSequences: boolean;
  readonly returnState: boolean;
  constructor(args: BidirectionalLayerArgs);
  build(inputShape: Shape | Shape[]): void;
  computeOutputShape(inputShape: Shape | Shape[]): Shape | Shape[];
  call(inputs: Tensor | Tensor[], kwargs: Kwargs): Tensor | Tensor[];
  resetStates(states?: Tensor | Tensor[]): void;
  get trainableWeights(): LayerVariable[];
  get nonTrainableWeights(): LayerVariable[];
}

Import

import { TimeDistributed, Bidirectional, Wrapper } from './layers/wrappers';

Layer Details

TimeDistributed

• Wraps any layer to apply it independently to each timestep.
• Input must be rank >= 3 (with a time dimension at axis 1).
• Uses rnn() internally: defines a step function that calls the wrapped layer on each timestep's data.
• Output shape inserts the timesteps dimension back: [batch, timesteps, ...layerOutputShape].

Bidirectional

Merge Mode	Output	Description
'concat' (default)	[batch, time?, units*2]	Concatenates forward and backward outputs
'sum'	[batch, time?, units]	Element-wise sum
'ave'	[batch, time?, units]	Element-wise average
'mul'	[batch, time?, units]	Element-wise product
null	[[forwardOutput], [backwardOutput]]	Returns separate outputs (no merge)

• Creates a forward copy and a backward copy (goBackwards=true) of the wrapped RNN layer.
• The returnSequences and returnState properties are inherited from the wrapped RNN layer.
• If returnState=true, forward and backward states are appended to the output.

I/O Contract

TimeDistributed

• Input: Rank-3+ tensor [batch, timesteps, ...].
• Output: [batch, timesteps, ...wrappedLayerOutput]. The wrapped layer sees input shape [batch, ...] (time dimension removed).
• Constraint: Input must be at least 3D.

Bidirectional

• Input: Rank-3 tensor [batch, timesteps, features] (standard RNN input).
• Output: Depends on mergeMode, returnSequences, and returnState (see table above).
• Constraint: The wrapped layer must be an RNN instance.

Usage Examples

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

// TimeDistributed: apply Dense to each frame in a video sequence
const model = tf.sequential();
model.add(tf.layers.timeDistributed({
  layer: tf.layers.dense({ units: 64 }),
  inputShape: [10, 128]  // 10 timesteps, 128 features
}));
// Output shape: [batch, 10, 64]

// TimeDistributed with Conv2D for video frames
const videoModel = tf.sequential();
videoModel.add(tf.layers.timeDistributed({
  layer: tf.layers.conv2d({ filters: 32, kernelSize: 3 }),
  inputShape: [10, 64, 64, 3]  // 10 frames of 64x64 RGB
}));

// Bidirectional LSTM
const biModel = tf.sequential();
biModel.add(tf.layers.bidirectional({
  layer: tf.layers.lstm({ units: 64, returnSequences: true }),
  inputShape: [100, 32],
  mergeMode: 'concat'
}));
// Output shape: [batch, 100, 128] (64*2 due to concat)

biModel.add(tf.layers.bidirectional({
  layer: tf.layers.lstm({ units: 32 }),
  mergeMode: 'sum'
}));
// Output shape: [batch, 32]

// Bidirectional with returnState for encoder
const encoder = tf.layers.bidirectional({
  layer: tf.layers.lstm({ units: 256, returnState: true }),
  mergeMode: 'concat'
});

Related Pages

• Tensorflow_Tfjs_Recurrent_Layers - RNN, LSTM, GRU layers wrapped by Bidirectional; rnn() function used by TimeDistributed
• Tensorflow_Tfjs_Exports_Layers - Public API factory functions tf.layers.timeDistributed() and tf.layers.bidirectional()
• Tensorflow_Tfjs_Merge_Layers - Merge operations conceptually similar to Bidirectional's merge modes
• Tensorflow_Tfjs_Layer_Variables - Weight delegation from wrapper to wrapped layer