ttnn.conv2d

ttnn.conv2d(input_tensor: ttnn.Tensor, weight_tensor: ttnn.Tensor, device: ttnn.MeshDevice, in_channels: int, out_channels: int, batch_size: int, input_height: int, input_width: int, kernel_size: tuple[int, int], stride: tuple[int, int], padding: tuple[int, int] or tuple[int, int, int, int], dilation: tuple[int, int], groups: int, *, dtype: ttnn.DataType | None, bias_tensor: ttnn.Tensor | None, conv_config: ttnn.Conv2dConfig | None, compute_config: ttnn.DeviceComputeKernelConfig | None, memory_config: ttnn.MemoryConfig | None, slice_config: ttnn.Conv2dSliceConfig | None, return_output_dim: bool | None, return_weights_and_bias: bool | None) → The output tensor, output height and width, and the preprocessed weights and bias.

Applies a 2D convolution over an input signal composed of several input planes.

Performs a 2D convolution between the input tensor and weight tensor. A 2D kernel (weights tensor) traverses the image (4D input tensor) and a dot product is computed over the overlapping region. For more information, refer to CNNs on Tenstorrent Architectures tech report.

Parameters:

• input_tensor (ttnn.Tensor) – The input tensor in [N, H, W, C] format. The tensor can be on either the host or the device.

• weight_tensor (ttnn.Tensor) – The convolution weights, typically in [out_channels, in_channels // groups, kernel_height, kernel_width] format.

• device (ttnn.MeshDevice) – This is a Tenstorrent-specific parameter. The device which will run the operation.

• in_channels (int) – Number of channels in the input tensor.

• out_channels (int) – Number of channels produced by the convolution.

• batch_size (int) – The batch size of the input tensor.

• input_height (int) – This is a Tenstorrent-specific parameter. The height of the input tensor.

• input_width (int) – This is a Tenstorrent-specific parameter. The width of the input tensor.

• kernel_size (tuple[int, int]) – The size of the convolving kernel.

• stride (tuple[int, int]) – The stride of the convolution. Default: (1, 1).

• padding (tuple[int, int] or tuple[int, int, int, int]) – Zero-padding added to both sides of the input. Default: (0, 0). [pad_height, pad_width] or [pad_top, pad_bottom, pad_left, pad_right].

• dilation (tuple[int, int]) – The spacing between kernel elements. Default: (1, 1).

• groups (int) – Number of blocked connections from input channels to output channels. Default: 1.

Keyword Arguments:

• dtype (ttnn.DataType, optional) – The data type of the output tensor. If not provided, it is inferred from the input tensor.

• bias_tensor (ttnn.Tensor, optional) – The bias tensor to be added. Default: None.

• conv_config (ttnn.Conv2dConfig, optional) – Configuration for convolution. Default: None.

• compute_config (ttnn.DeviceComputeKernelConfig, optional) – Configuration for compute kernel. Default: None

• memory_config (ttnn.MemoryConfig, optional) – Output Tensor’s Memory Configuration. Default: None.

• slice_config (ttnn.Conv2dSliceConfig, optional) – Configuration for slicing input & output tensors in DRAM. If set to None and input is in DRAM, DRAM slicing is automatically enabled. Default: None.

• return_output_dim (bool, optional) – If true, the op also returns the height and width of the output tensor in [N, H, W, C] format. Default: False

• return_weights_and_bias (bool, optional) – If true, the op also returns the preprocessed weight and bias on device. Default: False

Returns:

The output tensor, output height and width, and the preprocessed weights and bias.

• ttnn.Tensor: Default. The output tensor, when return_output_dim = False and return_weights_and_bias = False

• tuple[ttnn.Tensor, tuple[int, int]]: The output tensor, and its height and width, if return_output_dim = True

• tuple[ttnn.Tensor, tuple[ttnn.Tensor, ttnn.Tensor]]: The output tensor, and its height and width, if return_weights_and_bias = True

• tuple[ttnn.Tensor, tuple[int, int], tuple[ttnn.Tensor, ttnn.Tensor]]: The output tensor, and its height and width, if return_output_dim = True and return_weights_and_bias = True

Note

The input_tensor supports the following data type and layout:

input_tensor

dtype	layout
FLOAT32	ROW_MAJOR, TILE
BFLOAT16	ROW_MAJOR, TILE
BFLOAT8_B	TILE

The output_tensor supports the following data type and layout:

output_tensor

dtype	layout
FLOAT32	ROW_MAJOR, TILE
BFLOAT16	ROW_MAJOR, TILE
BFLOAT8_B	TILE

The weights_tensor on the host, supports the following data type and layout:

weights_tensor (host)

dtype	layout
FLOAT32	ROW_MAJOR
BFLOAT16	ROW_MAJOR

The weights_tensor prepared on device, supports the following data type and layout:

weights_tensor (prepared on device)

dtype	layout
FLOAT32	TILE
BFLOAT16	TILE
BFLOAT8_B	TILE

The bias_tensor on the host, supports the following data type and layout:

bias_tensor (host)

dtype	layout
FLOAT32	ROW_MAJOR
BFLOAT16	ROW_MAJOR

The bias_tensor prepared on device, supports the following data type and layout:

bias_tensor (prepared on device)

dtype	layout
FLOAT32	TILE
BFLOAT16	TILE
BFLOAT8_B	TILE