ttnn.embedding

ttnn.embedding(input_tensor: ttnn.Tensor, weight: ttnn.Tensor, *, memory_config: ttnn.MemoryConfig = input tensor memory config, output_tensor: ttnn.Tensor = None, padding_idx: int | None, layout: ttnn.Layout = ttnn.ROW_MAJOR_LAYOUT, embeddings_type: ttnn.EmbeddingsType = ttnn._ttnn.operations.embedding.EmbeddingsType.GENERIC, dtype: ttnn.DataType = None) → ttnn.Tensor

Retrieves word embeddings using input_tensor. The input_tensor is a list of indices, and the embedding matrix, and the output is the corresponding word embeddings.

Parameters:

  • input_tensor (ttnn.Tensor) – the input indices tensor.

  • weight (ttnn.Tensor) – the embeddings tensor that corresponds to the indices tensor.

Keyword Arguments:

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

  • output_tensor (ttnn.Tensor, optional) – Preallocated output tensor. Defaults to None.

  • padding_idx (int, optional) – the padding token. Default to None.

  • layout (ttnn.Layout) – the layout of the output tensor. Defaults to ttnn.ROW_MAJOR_LAYOUT.

  • embeddings_type (ttnn.EmbeddingsType) – the type of embeddings. Defaults to ttnn._ttnn.operations.embedding.EmbeddingsType.GENERIC.

  • dtype (ttnn.DataType, optional) – the data type for the output tensor. Defaults to None.

Returns:

ttnn.Tensor – the output tensor of layout == layout or layout of the weights tensor.

Example

# device_id = 0
# device = ttnn.open_device(device_id=device_id)

# Create a tensor containing indices into the embedding matrix
tensor = ttnn.to_device(
    ttnn.from_torch(torch.tensor([[1, 2, 4, 5], [4, 3, 2, 9]]), dtype=ttnn.uint32), device=device
)
# Create an embedding matrix containing 10 tensors of size 4
weight = ttnn.rand((10, 4), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=device)

# Perform embedding lookup
output = ttnn.embedding(tensor, weight)
logger.info(f"Output: {output}")