ttnn.gather

ttnn.gather(input: ttnn.Tensor, dim: int, index: ttnn.Tensor) → None

The gather operation extracts values from the input tensor based on indices provided in the index tensor along a specified dimension.

The input tensor and the index tensor must have the same number of dimensions. For all dimensions except the specified one (dim), the size of the index tensor must not exceed the size of the input tensor. The output tensor will have the same shape as the index tensor. Note that the input and index tensors do not broadcast against each other.

Parameters:

  • input (ttnn.Tensor) – The source tensor from which values are gathered.

  • dim (int) – The dimension along which values are gathered.

  • index (ttnn.Tensor) – A tensor containing the indices of elements to gather, with the same number of dimensions as the input tensor.

Keyword Arguments:

  • sparse_grad (bool, optional) – If True, the gradient computation will be sparse. Defaults to False.

  • memory_config (ttnn.MemoryConfig, optional) – Specifies the memory configuration for the output tensor. Defaults to None.

  • out (ttnn.Tensor, optional) – A preallocated tensor to store the gathered values. Defaults to None.

Additional Information:

  • Currently, the sparse_grad argument is not supported.

Note

Supported dtypes and layout for input tensor values:

Dtypes

Layouts

BFLOAT16, FLOAT32

TILE

UINT16, UINT32

TILE

INT32

TILE

Supported dtypes and layout for index tensor values:

Dtypes

Layouts

UINT16, UINT32

TILE
Memory Support:

  • Interleaved: DRAM and L1

Example

# Create a tensor and an index tensor
ttnn_input = ttnn.rand((4, 4), dtype=ttnn.bfloat16, layout=ttnn.Layout.TILE, device=device)
ttnn_index = ttnn.rand((4, 2), dtype=ttnn.uint16, layout=ttnn.Layout.TILE, device=device)

# Gather elements along dimension 1 using the index tensor
ttnn.gather(ttnn_input, 1, index=ttnn_index)