Embedding

Scenario Description

Map discrete IDs to dense vectors. Currently, KuDNN supports the torch.float16 data type. For other data types, see the open-source branch.

Sample Code

import torch
import torch.nn as nn

# Enable KuDNN.
torch._C._set_kdnn_enabled(True)

# Construct the Embedding layer.
embed = nn.Embedding(num_embeddings=1000, embedding_dim=128, dtype=torch.float16)

# Token index of the 2 × 2 input data
input_ids = torch.LongTensor([[1, 2], [3, 4]])
embeddings = embed (input_ids) # Output: [2, 2, 128]

print(embeddings)

Parent topic: Examples