ttnn.format_output_tensor

ttnn.format_output_tensor(output: ttnn._ttnn.tensor.Tensor, shape: list[int], device: ttnn._ttnn.multi_device.MeshDevice, target_layout: ttnn._ttnn.tensor.Layout, target_mem_config: ttnn._ttnn.tensor.MemoryConfig | None = None) → ttnn._ttnn.tensor.Tensor

Formats tensor to target layout and unpads to shape.

Parameters:

  • output (ttnn.Tensor) – Output tensor to format.

  • shape (ttnn.Shape) – Desired shape of the tensor.

  • device (ttnn.device.Device) – Device where the tensor will be moved.

  • target_layout (ttnn.Layout) – Desired tensor layout.

  • target_mem_config (ttnn.MemoryConfig, optional) – Desired memory config. Defaults to None.

Returns:

ttnn.Tensor – Formatted tensor.

Note

This functionality is planned for deprecation in the future.

Example

>>> # Assuming we have a padded tensor of shape [1, 2, 4, 4] with padding of [1, 1, 1, 1] of layout=ttnn.TILE_LAYOUT
>>> unpadded_tensor = ttnn.format_output_tensor(output_tensor, shape=[1, 2, 2, 2], device=device, target_layout=ttnn.ROW_MAJOR_LAYOUT, output_mem_config)