Implementation:Tensorflow Tfjs Embedding Layer

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

Overview

This module implements the Embedding layer for TensorFlow.js Layers, which turns positive integer indices into dense vectors of fixed size. It is commonly used as the first layer in NLP models to map word indices to embedding vectors. The layer supports masking (for variable-length sequences via maskZero), custom initializers, regularizers, and constraints on the embeddings matrix.

Code Reference

Source Location

tfjs-layers/src/layers/embeddings.ts (GitHub)

Key Imports

import {notEqual, reshape, serialization, Tensor, tidy, zerosLike} from '@tensorflow/tfjs-core';
import * as K from '../backend/tfjs_backend';
import {Layer, LayerArgs} from '../engine/topology';

Layer Class

Embedding

export class Embedding extends Layer {
  static className = 'Embedding';
  constructor(args: EmbeddingLayerArgs);
  public override build(inputShape: Shape | Shape[]): void;
  override computeMask(inputs: Tensor | Tensor[], mask?: Tensor | Tensor[]): Tensor;
  override computeOutputShape(inputShape: Shape | Shape[]): Shape | Shape[];
  override call(inputs: Tensor | Tensor[], kwargs: Kwargs): Tensor | Tensor[];
  override getConfig(): serialization.ConfigDict;
}

EmbeddingLayerArgs

export interface EmbeddingLayerArgs extends LayerArgs {
  inputDim: number;                    // vocabulary size (max integer index + 1)
  outputDim: number;                   // dimension of the embedding vectors
  embeddingsInitializer?: InitializerIdentifier | Initializer;  // default: 'randomUniform'
  embeddingsRegularizer?: RegularizerIdentifier | Regularizer;
  activityRegularizer?: RegularizerIdentifier | Regularizer;
  embeddingsConstraint?: ConstraintIdentifier | Constraint;
  maskZero?: boolean;                  // mask padding index 0
  inputLength?: number | number[];     // fixed input sequence length
}

Key Implementation Details

• build: Creates an embeddings weight matrix of shape [inputDim, outputDim].
• call: Casts input to int32, gathers rows from the embeddings matrix, and reshapes to the computed output shape.
• computeMask: When maskZero is true, returns a boolean tensor where non-zero input positions are true.
• computeOutputShape: Appends outputDim to the input shape. If inputLength is specified, validates and uses it for the sequence dimensions.

I/O Contract

Method	Input	Output
call	Integer tensor of shape [batch, ...sequenceDims]	Float tensor of shape [batch, ...sequenceDims, outputDim]
computeMask	Integer tensor	Boolean mask tensor (or null if maskZero is false)
computeOutputShape	Input shape array	Output shape with outputDim appended

Usage Example

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

const model = tf.sequential();
model.add(tf.layers.embedding({
  inputDim: 10000,    // vocabulary size
  outputDim: 128,     // embedding dimension
  inputLength: 50,    // sequence length
  maskZero: true
}));
// Input: [batch, 50] of integers in [0, 9999]
// Output: [batch, 50, 128] of float32

Related Pages

• Tensorflow_Tfjs_Constraints - Constraints for the embeddings matrix
• Tensorflow_Tfjs_Regularizers - Regularizers for the embeddings matrix
• Tensorflow_Tfjs_Exports_Initializers - Initializers for the embeddings matrix