ttnn.permute

ttnn.permute(input_tensor: ttnn.Tensor, dim: number, *, memory_config: ttnn.MemoryConfig = None, pad_value: float = 0.0) → List of ttnn.Tensor

Permutes the dimensions of the input tensor according to the specified permutation.

Parameters:

  • input_tensor (ttnn.Tensor) – the input tensor.

  • dim (number) – tthe permutation of the dimensions of the input tensor.

Keyword Arguments:

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

  • pad_value (float, optional) – padding value for when tiles are broken in a transpose. Defaults to 0.0.

Returns:

List of ttnn.Tensor – the output tensor.

Example

# Create a tensor to permute
input_tensor = ttnn.rand((1, 1, 64, 32), dtype=ttnn.bfloat16, device=device)

# Permute the tensor dimensions
permuted_tensor = ttnn.permute(input_tensor, (0, 1, 3, 2))
logger.info("Permuted Tensor Shape:", permuted_tensor.shape)  # Permuted Tensor Shape: Shape([1, 1, 32, 64])