ttnn.concat

ttnn.concat(input_tensor: List of ttnn.Tensor, dim: number, *, sub_core_grids: ttnn.CoreRangeSet = None) → ttnn.Tensor
Parameters:

  • input_tensor (List of ttnn.Tensor) – the input tensors.

  • dim (number) – the concatenating dimension.

Keyword Arguments:

  • memory_config (ttnn.MemoryConfig, optional) – Memory configuration for the operation. Defaults to None.

  • output_tensor (ttnn.Tensor, optional) – Preallocated output tensor. Defaults to None.

  • groups (int, optional) – When groups is set to a value greater than 1, the inputs are split into N groups partitions, and elements are interleaved from each group into the output tensor. Each group is processed independently, and elements from each group are concatenated in an alternating pattern based on the number of groups. This is useful for recombining grouped convolution outputs during residual concatenation. Defaults to 1. Currently, groups > 1 is only supported for two height sharded input tensors.

  • sub_core_grids (ttnn.CoreRangeSet, optional) – Sub-core grid to use for interleaved (L1 or DRAM) output tensors. If provided, the concatenation will run on the specified sub-core grid instead of the full compute grid. Defaults to None.

Returns:

ttnn.Tensor – the output tensor.

Example

# Create two tensors to concatenate
tensor1 = ttnn.rand([1, 1, 64, 32], dtype=ttnn.bfloat16, device=device)
tensor2 = ttnn.rand([1, 1, 64, 32], dtype=ttnn.bfloat16, device=device)

# Concatenate along dimension 3
concatenated_tensor = ttnn.concat([tensor1, tensor2], dim=3)
logger.info(
    "Concatenated Tensor Shape:", concatenated_tensor.shape
)  # Concatenated Tensor Shape: Shape([1, 1, 64, 64])