Implementation:Tensorflow Tfjs CPU Concat Kernel

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

Overview

The Concat kernel concatenates a list of tensors along a specified axis on the CPU backend. It handles edge cases such as empty output shapes, single non-empty inputs, and complex64 tensors. For the general case, the implementation reshapes all inputs into 2D matrices and delegates to concatImpl for efficient buffer concatenation.

Code Reference

Source Location

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

Signature

export function concat(
    args: {inputs: ConcatInputs, backend: MathBackendCPU, attrs: ConcatAttrs}):
    TensorInfo

Imports

import {backend_util, Concat, ConcatAttrs, ConcatInputs, KernelConfig, KernelFunc, TensorInfo, util} from '@tensorflow/tfjs-core';
import {MathBackendCPU} from '../backend_cpu';
import {complex} from './Complex';
import {concatImpl} from './Concat_impl';
import {identity} from './Identity';
import {imag} from './Imag';
import {real} from './Real';
import {reshape} from './Reshape';

Kernel Registration

export const concatConfig: KernelConfig = {
  kernelName: Concat,
  backendName: 'cpu',
  kernelFunc: concat as unknown as KernelFunc
};

I/O Contract

Parameter	Type	Description
inputs	ConcatInputs (array of TensorInfo)	The tensors to concatenate
attrs.axis	number	The axis along which to concatenate
Return	TensorInfo	A single tensor formed by concatenating the inputs

Implementation Details

1. The axis parameter is parsed and validated against input shapes using backend_util.assertParamsConsistent.
2. Empty output: If the total output size is zero, returns an empty tensor immediately.
3. Single input: If only one non-empty tensor remains after filtering, returns a copy via identity.
4. Complex64 path: Separates real and imaginary parts, concatenates each independently, and recombines via complex.
5. General path: Reshapes each input into a 2D tensor by splitting dimensions around the concat axis, calls concatImpl along axis 1, then reshapes back to the final output shape. Uses a simplyConcat flag when the first dimension is 1 for an optimized code path.

Usage Examples

import {concat} from './Concat';

// Concatenate two tensors along axis 0
const result = concat({
  inputs: [tensorA, tensorB],
  backend: cpuBackend,
  attrs: {axis: 0}
});

Related Pages

• Principle:Tensorflow_Tfjs_CPU_Kernel_Registration