This document is relevant for: Trn2, Trn3

Max Pooling 3D Backward Kernel API Reference#

Backward pass of separable 3D max pooling that uses the input and output of 3D max pooling forward and the backward gradient.

Background#

The max_pooling_3d_backward kernel computes the backward pass of 3D max pooling from the forward input, the forward output, and the output gradient.

API Reference#

Source code for this kernel API can be found at: max_pooling_3d_backward.py

max_pooling_3d_backward#

nkilib.experimental.pooling.max_pooling_3d_backward(grad_output: nl.NkiTensor, input_tensor: nl.NkiTensor, output_tensor: nl.NkiTensor, kernel_size: Union[int, tuple], stride: Optional[Union[int, tuple]] = None, padding: Union[int, tuple] = 0, data_format: str = 'NCDHW', output_format: Optional[str] = None) → nl.NkiTensor#

Backward pass of separable 3D max pooling that uses the input and output of 3D max pooling forward and the backward gradient.

Parameters:

  • grad_output (nl.NkiTensor) – Gradient w.r.t. the forward output, in HBM, output_format layout, forward-output shape.

  • input_tensor (nl.NkiTensor) – The forward input, data_format layout.

  • output_tensor (nl.NkiTensor) – The forward output (per-window maxima), output_format layout, same shape as grad_output.

  • kernel_size (Union[int, tuple]) – Pooling window (kD, kH, kW) (int broadcasts).

  • stride (Optional[Union[int, tuple]]) – Stride (sD, sH, sW); defaults to kernel_size.

  • padding (Union[int, tuple]) – Implicit -inf padding (pD, pH, pW).

  • data_format (str) – Layout of input_tensor and grad_input, “NCDHW” or “NDHWC”.

  • output_format (Optional[str]) – Layout of grad_output / output_tensor; defaults to data_format.

Returns:

Gradient w.r.t. the forward input, in HBM, data_format layout, same spatial extents as input_tensor.

Return type:

nl.NkiTensor

Dimensions:

  • B: batch.

This document is relevant for: Trn2, Trn3