This document is relevant for: Trn2, Trn3

Conv3D Transpose Kernel API Reference#

3D Transposed Convolution kernel.

Implements ConvTranspose3d operation by delegating to an embedded 3D convolution core with remapped parameters. The transposed convolution is achieved by mapping forward stride to input_dilation, remapping padding to dilation * (K - 1) - padding, and fixing kernel stride at 1.

Background#

The conv3d_transpose kernel implements a 3D transposed convolution (ConvTranspose3d).

API Reference#

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

conv3d_transpose#

nkilib.experimental.conv.conv3d_transpose(x_in: nl.NkiTensor, filters: nl.NkiTensor, bias: Optional[nl.NkiTensor] = None, stride: tuple[int, int, int] = (1, 1, 1), padding: tuple[int, int, int] = (0, 0, 0), dilation: tuple[int, int, int] = (1, 1, 1), activation_fn: Optional[ActFnType] = None, lnc_shard: bool = False, filter_shape: str = _FILTER_SHAPE_KDHW_CI_CO, sbm: Optional[SbufManager] = None, use_auto_allocation: bool = False) → nl.NkiTensor#

3D Transposed Convolution kernel.

Parameters:

  • x_in (nl.NkiTensor) – [B, C_in, D, H, W], Input tensor on HBM.

  • filters (nl.NkiTensor) – Filter weights on HBM with spatial axes flipped. Shape depends on filter_shape: - “KDHW_CI_CO” (default): [K_d, K_h, K_w, C_in, C_out] - “KDHW_CO_CI”: [K_d, K_h, K_w, C_out, C_in]

  • bias (Optional[nl.NkiTensor]) – [C_out], Optional bias tensor on HBM.

  • stride (tuple[int, int, int]) – (stride_d, stride_h, stride_w), Convolution strides.

  • padding (tuple[int, int, int]) – (pad_d, pad_h, pad_w), Padding for each spatial dimension.

  • dilation (tuple[int, int, int]) – (dilation_d, dilation_h, dilation_w), Filter dilation factors.

  • activation_fn (Optional[ActFnType]) – Optional activation function to apply after convolution.

  • lnc_shard (bool) – Enable LNC sharding across neuron cores.

  • filter_shape (str) – Storage layout of the filters tensor. Default is “KDHW_CI_CO”.

  • sbm (Optional[SbufManager]) – Optional caller-provided SBUF manager. When None, the kernel creates its own SbufManager.

  • use_auto_allocation (bool) – Must equal sbm.is_auto_alloc() when sbm is provided. When sbm is None this flag is unused.

Returns:

[B, C_out, D_out, H_out, W_out], Output tensor on HBM.

Return type:

nl.ndarray

Dimensions:

  • B: Batch size

  • C_in: Number of input channels

  • C_out: Number of output channels

  • D: Input depth dimension

  • H: Input height dimension

  • W: Input width dimension

  • K_d: Filter kernel depth

  • K_h: Filter kernel height

  • K_w: Filter kernel width

  • D_out: Output depth = (D - 1) * stride_d + dilation_d * (K_d - 1) - 2 * pad_d + 1

  • H_out: Output height = (H - 1) * stride_h + dilation_h * (K_h - 1) - 2 * pad_h + 1

This document is relevant for: Trn2, Trn3