Implementation:Tensorflow Tfjs CPU Einsum Kernel

Knowledge Sources	Tensorflow_Tfjs
Domains	Deep_Learning, CPU_Backend
Last Updated	2026-02-10 06:00 GMT

Overview

The Einsum kernel implements the Einstein summation convention on the CPU backend. It evaluates an arbitrary tensor contraction specified by a string equation (e.g., 'ij,jk->ik' for matrix multiplication). The implementation decomposes the equation into a sequence of transpose, reshape, multiply, and sum operations, following an optimized compute path determined by backend_util.getEinsumComputePath.

Code Reference

Source Location

tfjs-backend-cpu/src/kernels/Einsum.ts (GitHub)

Signature

export function einsum(
    args: {inputs: EinsumInputs, backend: MathBackendCPU, attrs: EinsumAttrs}):
    TensorInfo

Imports

import {backend_util, Einsum, EinsumAttrs, EinsumInputs, KernelConfig, KernelFunc, Tensor, TensorInfo, util} from '@tensorflow/tfjs-core';
import {MathBackendCPU} from '../backend_cpu';
import {multiply} from './Multiply';
import {reshape} from './Reshape';
import {sum} from './Sum';
import {transpose} from './Transpose';

Kernel Registration

export const einsumConfig: KernelConfig = {
  kernelName: Einsum,
  backendName: 'cpu',
  kernelFunc: einsum as unknown as KernelFunc
};

I/O Contract

Parameter	Type	Description
inputs	EinsumInputs (array of Tensor)	Input tensors to contract
attrs.equation	string	Einstein summation equation (e.g., 'ij,jk->ik')
Return	TensorInfo	Result tensor of the contraction

Implementation Details

1. Equation parsing: backend_util.decodeEinsumEquation extracts all dimension labels, summed dimensions, and per-input dimension identifiers.
2. Validation: backend_util.checkEinsumDimSizes ensures dimension sizes are consistent across inputs.
3. Compute path: backend_util.getEinsumComputePath determines the optimal order of contractions, returning a path (which dimensions to sum at each step) and steps (which input terms to process).
4. Step execution: For each step:
   • Each input term is optionally transposed (via transpose) and reshaped (via reshape) to align dimensions.
   • If an identity permutation is detected, the transpose is skipped.
   • Aligned tensors are element-wise multiplied (via multiply).
   • After processing all terms in a step, a sum reduction is applied along the path-specified axis (if applicable), reducing numDimsRemaining.
5. Cleanup: All intermediate tensors are disposed except the final output.

Usage Examples

import {einsum} from './Einsum';

// Matrix multiplication via einsum
const result = einsum({
  inputs: [matA, matB],
  backend: cpuBackend,
  attrs: {equation: 'ij,jk->ik'}
});

// Batch matrix multiply
const batchResult = einsum({
  inputs: [batchA, batchB],
  backend: cpuBackend,
  attrs: {equation: 'bij,bjk->bik'}
});

Related Pages

• Principle:Tensorflow_Tfjs_CPU_Kernel_Registration