Function maxPoolGrad

tfjs-backend-cpu/src/kernels/MaxPoolGrad.ts:23–95

(args: {
  inputs: MaxPoolGradInputs,
  backend: MathBackendCPU,
  attrs: MaxPoolGradAttrs
})

import {maxPoolPositions} from '../utils/pool_utils';

export function maxPoolGrad(args: {
  inputs: MaxPoolGradInputs,
  backend: MathBackendCPU,
  attrs: MaxPoolGradAttrs
}): TensorInfo {
  const {inputs, backend, attrs} = args;
  const {dy, input, output} = inputs;
  const x = input;
  assertNotComplex([input, output], 'maxPoolGrad');
  const {filterSize, strides, pad, dimRoundingMode} = attrs;

  const convInfo = backend_util.computePool2DInfo(
      x.shape as [number, number, number, number], filterSize, strides,
      1 /* dilations */, pad, dimRoundingMode);
  const xValues = backend.data.get(x.dataId).values as TypedArray;
  const maxPosBuf = buffer(
      convInfo.outShape, x.dtype,
      maxPoolPositions(xValues, x.shape, x.dtype, convInfo).values);
  const strideHeight = convInfo.strideHeight;
  const strideWidth = convInfo.strideWidth;
  const dilationHeight = convInfo.dilationHeight;
  const dilationWidth = convInfo.dilationWidth;
  const effectiveFilterHeight = convInfo.effectiveFilterHeight;
  const effectiveFilterWidth = convInfo.effectiveFilterWidth;
  const padLeft = effectiveFilterWidth - 1 - convInfo.padInfo.left;
  const padTop = effectiveFilterHeight - 1 - convInfo.padInfo.top;
  const dx =
      buffer<Rank.R4>(x.shape as [number, number, number, number], 'float32');

  const dyData = backend.data.get(dy.dataId).values as Float32Array;
  const dyBuf = buffer<Rank.R4>(
      dy.shape as [number, number, number, number], 'float32', dyData);

  for (let b = 0; b < convInfo.batchSize; ++b) {
    for (let d = 0; d < convInfo.inChannels; ++d) {
      for (let dxR = 0; dxR < convInfo.inHeight; ++dxR) {
        for (let dxC = 0; dxC < convInfo.inWidth; ++dxC) {
          // Shader code begins.
          const dyRCorner = dxR - padTop;
          const dyCCorner = dxC - padLeft;
          let dotProd = 0;
          for (let wR = 0; wR < effectiveFilterHeight; wR += dilationHeight) {
            const dyR = (dyRCorner + wR) / strideHeight;
            if (dyR < 0 || dyR >= convInfo.outHeight ||
                Math.floor(dyR) !== dyR) {
              continue;
            }
            for (let wC = 0; wC < effectiveFilterWidth; wC += dilationWidth) {
              const dyC = (dyCCorner + wC) / strideWidth;
              if (dyC < 0 || dyC >= convInfo.outWidth ||
                  Math.floor(dyC) !== dyC) {
                continue;
              }
              const maxPos = effectiveFilterHeight * effectiveFilterWidth - 1 -
                  (maxPosBuf.get(b, dyR, dyC, d) as number);
              const curPos = wR * effectiveFilterWidth + wC;

              const mask = maxPos === curPos ? 1 : 0;