ttnn.sharded_to_interleaved_partial

ttnn.sharded_to_interleaved_partial(input_tensor: ttnn.Tensor, cache_tensor: ttnn.Tensor, num_slices: int, slice_index: int, *, memory_config: ttnn.MemoryConfig | None = None, output_dtype: ttnn.DataType | None = None) → ttnn.Tensor

Converts a partial tensor from sharded_to_interleaved memory layout

Parameters:

  • input_tensor (ttnn.Tensor) – input tensor

  • cache_tensor (ttnn.Tensor) – cache tensor

  • num_slices (int) – Number of slices.

  • slice_index (int) – Slice index.

Keyword Arguments:

  • memory_config (Optional[ttnn.MemoryConfig]) – Memory configuration for the operation. Defaults to None.

  • output_dtype (Optional[ttnn.DataType]) – Output data type, defaults to same as input. Defaults to None.

Returns:

ttnn.Tensor – the cache tensor with the partial data written.

Example

>>> interleaved_tensor = ttnn.sharded_to_interleaved_partial(tensor, cache_tensor, 4, 2)