This document is relevant for: Trn2, Trn3

Average Pooling 3D Backward Kernel API Reference#

Backward pass of separable 3D average pooling.

Spreads each output gradient back over the input window it was averaged from, applying the 1/prod(kernel_size) scale, using a separable per-axis spread across D, H and W.

Background#

The avg_pooling_3d_backward kernel computes the backward pass of 3D average pooling. It spreads each output gradient back over its input window with the 1/prod(kernel_size) scale, one axis at a time across D, H and W.

API Reference#

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

avg_pooling_3d_backward#

nkilib.experimental.pooling.avg_pooling_3d_backward(grad_output: nl.NkiTensor, input_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 average pooling.

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. Only its shape and dtype are read an average pool’s gradient does not depend on the forward values.

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

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

  • padding (Union[int, tuple]) – Implicit zero 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; 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 size

  • C: Number of channels

  • D: Input depth (spatial)

  • H: Input height (spatial)

  • W: Input width (spatial)

  • Do: Output depth (pooled)

  • Ho: Output height (pooled)

This document is relevant for: Trn2, Trn3