Implementation:Tensorflow Tfjs CPU MaxPoolGrad Kernel

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

Overview

The MaxPoolGrad kernel computes the gradient of the max pooling operation for the CPU backend. It back-propagates the upstream gradient (dy) to the input tensor by determining which input positions contributed the maximum value in each pooling window, then routing the gradient only to those positions.

Code Reference

Source Location

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

Signature

export function maxPoolGrad(args: {
  inputs: MaxPoolGradInputs,
  backend: MathBackendCPU,
  attrs: MaxPoolGradAttrs
}): TensorInfo

Imports

import {backend_util, buffer, KernelConfig, KernelFunc, MaxPoolGrad, MaxPoolGradAttrs, MaxPoolGradInputs, Rank, TensorInfo, TypedArray} from '@tensorflow/tfjs-core';
import {MathBackendCPU} from '../backend_cpu';
import {assertNotComplex} from '../cpu_util';
import {maxPoolPositions} from '../utils/pool_utils';

Kernel Config

export const maxPoolGradConfig: KernelConfig = {
  kernelName: MaxPoolGrad,
  backendName: 'cpu',
  kernelFunc: maxPoolGrad as unknown as KernelFunc
};

I/O Contract

Inputs

Name	Type	Description
dy	TensorInfo	Upstream gradient from the output of max pool (4D: [batch, outHeight, outWidth, channels])
input	TensorInfo	Original input tensor to the forward max pool (4D: [batch, inHeight, inWidth, channels])
output	TensorInfo	Output of the forward max pool operation

Attributes

Name	Type	Description
filterSize	number	Height and width of the pooling filter
strides	number	Stride of the pooling window
pad	'same' | number	Padding mode or explicit pad amount
dimRoundingMode	'round' | 'ceil'	Rounding mode for output dimensions

Output

Returns a TensorInfo with float32 dtype and the same shape as the input tensor, containing the back-propagated gradients.

Algorithm

1. Computes pool2D convolution info using backend_util.computePool2DInfo.
2. Determines the position of the max value in each pooling window via maxPoolPositions.
3. Iterates over every input position [b, dxR, dxC, d].
4. For each input position, iterates over the effective filter window to find which output positions could map back to this input.
5. Compares the max position index to the current filter position; if they match (mask = 1), adds the corresponding dy value to the gradient accumulator.
6. Sets the accumulated gradient value in the output buffer at [b, dxR, dxC, d].

Usage Example

This kernel is invoked automatically by TensorFlow.js when computing gradients through a tf.maxPool operation:

import * as tf from '@tensorflow/tfjs-core';
import '@tensorflow/tfjs-backend-cpu';

const x = tf.tensor4d([1, 2, 3, 4, 5, 6, 7, 8, 9], [1, 3, 3, 1]);
const grads = tf.grads((input) => tf.maxPool(input, 2, 1, 'valid'));
const dx = grads([x]);

Related Pages

• Principle:Tensorflow_Tfjs_CPU_Kernel_Registration