Implementation:Tensorflow Tfjs Normalization Test

Knowledge Sources	Tensorflow_Tfjs
Domains	Testing, Layers_API
Last Updated	2026-02-10 06:00 GMT

Overview

This test suite validates normalization layers in TensorFlow.js Layers, covering both low-level normalization functions and the BatchNormalization and LayerNormalization layer APIs. The tests verify batch normalization in training mode (computing running mean/variance), inference mode (using stored statistics), and the underlying normalizeBatchInTraining and batchNormalization functions. Reference values are generated from Python Keras for numerical verification. The suite covers 2D, 3D, and 4D inputs, various axis configurations, and training integration with model.fit().

Code Reference

Source Location: tfjs-layers/src/layers/normalization_test.ts (1012 lines)

Repository: GitHub

Test Describe Blocks

normalizeBatchInTraining - Low-level function testing:
- 2D inputs, no broadcasting
- 3D inputs, no broadcasting
- 3D inputs, with broadcasting
batchNormalization - Low-level batch normalization function:
- 2D, 3D, and 4D inputs
- With and without gamma/beta
- Different axis configurations
BatchNormalization Layers: Symbolic - Shape inference for batch norm layers
BatchNormalization Layers: Tensor - Numerical correctness:
- Default axis, custom axis
- Training vs. inference mode
- Running mean/variance updates
- Integration with model.fit()
- Regularizer support
- Fused batch normalization
- Serialization (getConfig, fromConfig)
LayerNormalization Layer: Symbolic - LayerNorm shape inference
LayerNormalization Layer: Tensor - LayerNorm numerical correctness

I/O Contract

Inputs to tests:

2D tensors ([batch, features]), 3D tensors ([batch, time, features]), 4D tensors ([batch, height, width, channels])
Gamma (scale) and beta (offset) parameters
Reduction axes configuration
Training flag for controlling batch norm behavior

Expected outputs/assertions:

Normalized output values match Python Keras reference values
Running mean and variance are updated correctly during training
Inference uses stored running statistics, not batch statistics
Default epsilon value prevents division by zero
LayerNormalization normalizes along specified axes
Serialization preserves all configuration including axis, momentum, epsilon

Usage Example

describeMathCPUAndGPU('normalizeBatchInTraining', () => {
  it('2D, no broadcasting', () => {
    const x = tensor2d([[1, 2, 3, 4], [2, 4, 6, 8], [12, 11, 10, 9]], [3, 4]);
    const gamma = tensor1d([1, 1, 1, 1]);
    const beta = tensor1d([0, 0, 0, 0]);
    const reductionAxes = [0];
    const [normed, mean, variance] =
        normalizeBatchInTraining(x, gamma, beta, reductionAxes);
    expectTensorsClose(mean, tensor1d([5.0, 5.6666665, 6.3333335, 7.0]));
    expectTensorsClose(
        variance, tensor1d([24.666666, 14.888889, 8.222222, 4.6666665]));
  });
});

Test Coverage Summary

Category	Count	Details
normalizeBatchInTraining	3	2D/3D no broadcast, 3D with broadcast
batchNormalization function	10+	2D/3D/4D, with/without gamma/beta
BatchNormalization Layer	25+	Symbolic, tensor, training/inference, fit(), serialization
LayerNormalization	10+	Symbolic shape inference and tensor correctness
Test Environment	Mixed	CPU, GPU, WebGL2

Related Pages

Page Connections

Double-click a node to navigate. Hold to expand connections.

Principle

Implementation

Heuristic

Environment