ttnn.prepare_conv_bias

ttnn.prepare_conv_bias(*args, **kwargs)

TTNN Conv2D applies preprocessing to the bias tensors before performing the convolution operation, to convert the bias into a format suitable for the operation. This can be applied just once to the weights and bias tensors, and the resulting tensors can be reused for multiple invocations of the same convolution operation. The exact format of the weights and bias tensors depends on the input tensor parameters and the sharding scheme.

Parameters:

  • bias (ttnn.Tensor) – the bias tensor in PyTorch Conv2d format.

  • input_memory_config (ttnn.MemoryConfig) – the memory configuration for the input tensor.

  • input_layout (ttnn.Tensor) – the layout of the input tensor.

  • int – in_channels: number of input channels.

  • int – out_channels: number of output channels.

  • int – batch_size: batch size.

  • int – input_height: height of the input tensor.

  • int – input_width: width of the input tensor.

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

  • stride (tuple[int, int]) – stride of the cross-correlation.

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

  • dilation (tuple[int, int]) – spacing between kernel elements.

  • device (ttnn.IDevice) – the device to use.

  • groups (int) – number of blocked connections from input channels to output channels.

  • conv_config (ttnn.Conv2dConfig, None) – configuration for convolution. This config must have weights_dtype set to the same dtype as the processed weights tensor. Default: None

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

Returns:

The preprocessed bias tensor on device

Return type:

[ttnn.Tensor]: The preprocessed bias tensor on device