ttnn.slice

ttnn.slice(input_tensor: ttnn.Tensor, slice_start: List[int], slice_end: List[int], slice_step: List[int] = ``None`` (step = 1 for all dims, *, memory_config: ttnn.MemoryConfig = the input tensor's memory config, output_tensor: ttnn.Tensor = ``None``, pad_value: float | None, sub_core_grids: ttnn.CoreRangeSet = None) → ttnn.Tensor

Returns a sliced tensor. If the input tensor is on host, the slice will be performed on host, and if its on device it will be performed on device.

Parameters:

  • input_tensor (ttnn.Tensor) – Input tensor.

  • slice_start (List[int]) – Start indices of input tensor. Values along each dim must be in [0, input_tensor_shape[i]).

  • slice_end (List[int]) – End indices of input tensor (exclusive). Values along each dim must be in (0, input_tensor_shape[i]].

  • slice_step (List[int], optional) – Step size for each dim. Defaults to None (step = 1 for all dims).

Keyword Arguments:

  • memory_config (ttnn.MemoryConfig, optional) – Memory configuration for the output tensor. Defaults to the input tensor’s memory config.

  • output_tensor (ttnn.Tensor, optional) – Pre-allocated output tensor. Its shape must match the slice output. Defaults to None.

  • pad_value (float, optional) – Fill value for implicit tile padding on tiled tensors. Padding is undefined by default.

  • sub_core_grids (ttnn.CoreRangeSet, optional) – sub core grids for the operation. Defaults to None.

Note

Strided slicing (slice_step != 1) is not supported for bfloat8_b tensors.

Returns:

ttnn.Tensor – the output tensor.

Example

# Create a tensor to slice
input_tensor = ttnn.rand((1, 1, 64, 32), dtype=ttnn.bfloat16, layout=ttnn.Layout.TILE, device=device)

# Slice the tensor
sliced_tensor = ttnn.slice(input_tensor, [0, 0, 0, 0], [1, 1, 64, 16], [1, 1, 2, 1])
logger.info("Sliced Tensor Shape:", sliced_tensor.shape)  # Sliced Tensor Shape: Shape([1, 1, 32, 16])

# Create a tensor to slice without step
input_tensor = ttnn.rand((1, 1, 64, 32), dtype=ttnn.bfloat16, layout=ttnn.Layout.TILE, device=device)
output = ttnn.slice(input_tensor, [0, 0, 0, 0], [1, 1, 32, 32])
logger.info("Sliced Tensor Shape:", output.shape)  # Sliced Tensor Shape: Shape([1, 1, 32, 32])