Implementation:Tensorflow Tfjs CPU Conv3D Kernel

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

Overview

The Conv3D kernel implements a 3D convolution on the CPU backend. It convolves a 5D input tensor ([batch, depth, height, width, channels]) with a 5D filter, supporting configurable strides, dilations, and padding. The implementation uses a direct nested-loop approach with precomputed strides for memory offset calculation.

Code Reference

Source Location

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

Signature

export function conv3D(
    args: {inputs: Conv3DInputs, backend: MathBackendCPU, attrs: Conv3DAttrs}):
    TensorInfo

Imports

import {backend_util, Conv3D, Conv3DAttrs, Conv3DInputs, KernelConfig, KernelFunc, TensorBuffer, TensorInfo, TypedArray, util} from '@tensorflow/tfjs-core';
import {MathBackendCPU} from '../backend_cpu';
import {assertNotComplex} from '../cpu_util';

Kernel Registration

export const conv3DConfig: KernelConfig = {
  kernelName: Conv3D,
  backendName: 'cpu',
  kernelFunc: conv3D as unknown as KernelFunc
};

I/O Contract

Parameter	Type	Description
inputs.x	TensorInfo	5D input tensor of shape [batch, inDepth, inHeight, inWidth, inChannels]
inputs.filter	TensorInfo	5D filter tensor of shape [filterDepth, filterHeight, filterWidth, inChannels, outChannels]
attrs.strides	[number, number, number]	Stride for each spatial dimension
attrs.pad	'same' | number	Padding mode
attrs.dilations	[number, number, number]	Dilation rate for each spatial dimension
Return	TensorInfo	Output tensor of shape [batch, outDepth, outHeight, outWidth, outChannels]

Implementation Details

• Asserts that inputs are not complex tensors.
• Calls backend_util.computeConv3DInfo to derive output shape, padding, and stride metadata.
• Allocates a TensorBuffer for the output with the computed outShape.
• Uses eight nested loops: batch, output depth, filter depth, output height, filter height, output width, filter width, and input/output channels.
• Skips out-of-bounds input positions (padding regions) via boundary checks against inDepth, inHeight, and inWidth.
• Dilation is applied by multiplying the filter index by the dilation factor when computing input positions.
• Accumulates the convolution sum yVals[yOffset + d2] += xVal * wVals[wOffset + d2] for each output channel.

Usage Examples

import {conv3D} from './Conv3D';

const result = conv3D({
  inputs: {x: input5D, filter: filter5D},
  backend: cpuBackend,
  attrs: {strides: [1, 1, 1], pad: 'same', dilations: [1, 1, 1]}
});

Related Pages

• Principle:Tensorflow_Tfjs_CPU_Kernel_Registration