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.Tensorthe 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}")