This document is relevant for: Trn2, Trn3

Average Pooling 3D Kernel API Reference#

3D average pool kernel over the (D, H, W) axes. Pooling runs independently per (batch, channel).

Background#

The avg_pooling_3d kernel computes 3D average pooling over the (D, H, W) axes, independently for each (batch, channel) pair. It supports channels-first (NCDHW) and channels-last (NDHWC) layouts.

API Reference#

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

avg_pooling_3d#

nkilib.experimental.pooling.avg_pooling_3d(src_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#

3D average pool kernel over the (D, H, W) axes. Pooling runs independently per (batch, channel).

Parameters:

  • src_tensor (nl.NkiTensor) – Input in HBM. Shape is (B, C, D, H, W) for data_format=”NCDHW” (channels-first) or (B, D, H, W, C) for data_format=”NDHWC” (channels-last).

  • 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 zero padding (pD, pH, pW).

  • data_format (str) – Input layout, “NCDHW” (default) or “NDHWC”.

  • output_format (Optional[str]) – Output layout, “NCDHW” or “NDHWC”; defaults to data_format (output matches input).

Returns:

Pooled result in HBM, in output_format layout: (B, C, D_out, H_out, W_out) or (B, D_out, H_out, W_out, C) with out[a] = (in[a] + 2*pad[a] - kernel[a]) // stride[a] + 1.

Return type:

nl.NkiTensor

Dimensions:

  • B: batch.

This document is relevant for: Trn2, Trn3