ttnn.scatter_add

ttnn.scatter_add(input: ttnn.Tensor, dim: int, index: ttnn.Tensor, src: ttnn.Tensor, *, memory_config: ttnn.MemoryConfig = None, sub_core_grids: ttnn.CoreRangeSet = None) → ttnn.Tensor

Scatters the source tensor’s values along a given dimension according to the index tensor, adding source values associated with according repeated indices.

Parameters:

  • input (ttnn.Tensor) – the input tensor to scatter values onto.

  • dim (int) – the dimension to scatter along.

  • index (ttnn.Tensor) – the tensor specifying indices where values from the source tensor must go to.

  • src (ttnn.Tensor) – the tensor containing the source values to be scattered onto input.

Keyword Arguments:

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

  • sub_core_grids (ttnn.CoreRangeSet, optional) – specifies which cores scatter should run on. Defaults to None.

Returns:

ttnn.Tensor – the output tensor with scattered values.

Note

  • Input tensors must be interleaved and on device.

Example

# Create input, index, and source tensors for scatter_add
input_torch = torch.tensor([[1, 2], [3, 4]], dtype=torch.float32)
index_torch = torch.tensor([[0, 1], [1, 0]], dtype=torch.int64)
source_torch = torch.tensor([[5, 6], [7, 8]], dtype=torch.float32)

input_ttnn = ttnn.from_torch(input_torch, dtype=ttnn.float32, device=device, layout=ttnn.ROW_MAJOR_LAYOUT)
index_ttnn = ttnn.from_torch(index_torch, dtype=ttnn.int32, device=device, layout=ttnn.ROW_MAJOR_LAYOUT)
source_ttnn = ttnn.from_torch(source_torch, dtype=ttnn.float32, device=device, layout=ttnn.ROW_MAJOR_LAYOUT)
dim = 1

# Perform scatter_add operation (adds source values to input at specified indices)
output = ttnn.scatter_add(input_ttnn, dim, index_ttnn, source_ttnn)
logger.info(f"Scatter add operation result: {output}")