Implementation:Tensorflow Tfjs EinsumDense Test

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

Overview

This test suite validates the EinsumDense layer, which uses Einstein summation notation to define flexible dense (fully-connected) operations. This layer is a key component in transformer architectures, enabling arbitrary tensor contractions specified by an equation string. The tests verify correct weight shape inference, bias shape inference, output shape computation, and the analyzeEinsumString utility function across a variety of equation patterns including 1D, 2D, 3D, and BERT-style attention projections.

Code Reference

Source Location: tfjs-layers/src/layers/nlp/einsum_dense_test.ts (317 lines)

Repository: GitHub

Test Describe Blocks

EinsumDense - Comprehensive tests organized by equation pattern:
- _1d_end_weight: equation ab,b->a (dot product)
- _2d_middle_weight: equation ab,bc->ac (standard dense)
- _3d_bert: equation abc,cde->abde (BERT-style projection)
- _3d_3_bias: equation with bias on axis 'e'
- _3d_2_bias: equation with bias on axis 'd'
- _3d_1_3_bias: equation with bias on axes 'be'
- Additional test cases for ellipsis notation and various axis combinations
- analyzeEinsumString tests for parsing equation strings

I/O Contract

Inputs to tests:

Einsum equation strings (e.g., 'ab,bc->ac', 'abc,cde->abde')
Bias axes specifications (e.g., 'e', 'd', 'be', or null)
Input shapes and output shapes
Concrete tensors for forward pass verification

Expected outputs/assertions:

Weight shapes correctly inferred from equation (e.g., [32, 64] for ab,bc->ac with input [null, 32] and output [64])
Bias shapes correctly inferred from bias axes (e.g., [4] for bias axis 'e' with output dim 4)
Output shapes correctly computed (e.g., [null, 1, 3, 4] for BERT-style projection)
analyzeEinsumString correctly parses equations into input/output labels

Usage Example

describe('EinsumDense', () => {
  const combinations: EinsumDenseTestCaseArgs[] = [
    {
      testcaseName: '_2d_middle_weight',
      equation: 'ab,bc->ac',
      biasAxes: null,
      inputShape: [null, 32],
      outputShape: [64],
      expectedWeightShape: [32, 64],
      expectedBiasShape: null,
      expectedOutputShape: [null, 64],
    },
  ];

  for (const combo of combinations) {
    it(`test${combo.testcaseName}`, () => {
      const layer = new EinsumDense({
        equation: combo.equation,
        outputShape: combo.outputShape,
        biasAxes: combo.biasAxes,
      });
      const symbolicInput = input({shape: combo.inputShape.slice(1)});
      layer.apply(symbolicInput);
      expect(layer.kernel.shape).toEqual(combo.expectedWeightShape);
    });
  }
});

Test Coverage Summary

Category	Count	Details
Equation Patterns	10+	1D, 2D, 3D, BERT-style, ellipsis notation
Weight Shape Tests	10+	Shape inference from equation for each pattern
Bias Shape Tests	5+	Single axis, multiple axes, null bias
Output Shape Tests	10+	End-to-end shape computation
String Analysis	3+	analyzeEinsumString parsing
Test Environment	Standard	`describe`

Related Pages

Page Connections

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

Principle

Implementation

Heuristic

Environment