For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Python class

Embedding

Embedding

class max.nn.Embedding(vocab_size, hidden_dim, dtype, device, quantization_encoding=None, name=None)

source

Bases: Module

A lookup table for embedding integer indices into dense vectors.

When called, Embedding maps each integer index to a dense vector of fixed size. It accepts a TensorValueLike of integer indices with shape (batch, ..., num_indices) and returns a TensorValue of shape (batch, ..., num_indices, hidden_dim) containing the corresponding embedding vectors.

Embedding weights are stored on the CPU but are moved to the specified device during model initialization.

from max.driver import Accelerator, CPU, accelerator_count
from max.dtype import DType
from max.graph import DeviceRef, Graph, TensorType
from max.nn import Embedding

device = Accelerator() if accelerator_count() > 0 else CPU()
device_ref = DeviceRef.from_device(device)

embedding_layer = Embedding(
    vocab_size=1000,
    hidden_dim=256,
    dtype=DType.float32,
    device=device_ref,
    name="embeddings",
)

indices_type = TensorType(DType.uint32, [1, 8], device=device_ref)
with Graph("embedding", input_types=[indices_type]) as graph:
    embeddings = embedding_layer(graph.inputs[0])
    graph.output(embeddings)

Initializes the embedding layer with the given arguments.

Parameters:

• vocab_size (int) – The number of unique items in the vocabulary. Indices must be in the range [0, vocab_size).
• hidden_dim (int) – The dimensionality of each embedding vector.
• dtype (DType) – The data type of the embedding weights.
• device (DeviceRef) – The device where embedding lookups are executed. Model init transfers the initially CPU-resident weights to this device.
• quantization_encoding (QuantizationEncoding | None) – Optional quantization encoding for the weights.
• name (str | None) – The name identifier for the embedding weight matrix.

device

device: DeviceRef

source

The device on which embedding lookup is performed.

weight

weight: Weight

source

The embedding weight matrix stored on the CPU. Model init moves weights to the device specified in device.