ttnn.index_fill

ttnn.index_fill(input: ttnn.Tensor, dim: int, index: ttnn.Tensor, value: int or float, *, memory_config: ttnn.MemoryConfig = None) → ttnn.Tensor

Fills the input tensor with the given value at the specified indices along the specified dimension.

Parameters:

  • input (ttnn.Tensor)The input tensor.

  • dim (int)The dimension along which to fill the value.

  • index (ttnn.Tensor)A tensor containing the indices along dim to fill with the given value.

  • value (int or float)The value which will be used to fill the output tensor.

Keyword Arguments:

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

Returns:

ttnn.TensorThe output tensor.

Note

This operation supports tensors according to the following data types and layouts:

input tensor

dtype - layout

BFLOAT16, FLOAT32, INT32 - ROW_MAJOR
index tensor

dtype - layout

UINT32 - ROW_MAJOR

UINT32 - TILE
Memory Support:

  • Interleaved (DRAM or L1): output may be any supported layout

  • Height sharded (DRAM or L1): output may be any supported layout

  • Width sharded (DRAM or L1): output may be interleaved, height sharded, width sharded, or block sharded; for width/block sharded output the column shard width (shard_spec.shape[1]) must match the input

  • Block sharded (DRAM or L1): output may be interleaved, height sharded, block sharded, or width sharded; for width/block sharded output the column shard width (shard_spec.shape[1]) must match the input

Limitations:

  • The input tensor must be on the device.

  • The index tensor must be on the device and must be a 1D tensor.

  • The dim must be less than the number of dimensions of the input tensor and >= 0.

  • The value must be a float or int and must match the dtype of the input tensor.

Example

# Create a TT-NN tensor with values filled at the specified indices along the specified dimension
tt_input = ttnn.rand([32, 32], dtype=ttnn.bfloat16, layout=ttnn.ROW_MAJOR_LAYOUT, device=device)
tt_index = ttnn.Tensor([0, 31], [2], ttnn.uint32, ttnn.ROW_MAJOR_LAYOUT, device)

output = ttnn.index_fill(
    tt_input, 1, tt_index, 10.0
)  # Need to ensure 10.0 is a float to match the bfloat16 dtype of the input tensor
logger.info("TT-NN index_fill tensor:", output)