mlx.nn.Embedding

class Embedding(num_embeddings: int, dims: int)

Implements a simple lookup table that maps each input integer to a high-dimensional vector.

Typically used to embed discrete tokens for processing by neural networks.

Parameters:

• num_embeddings (int) – How many possible discrete tokens can we embed. Usually called the vocabulary size.

• dims (int) – The dimensionality of the embeddings.

Methods

as_linear(x)	Call the embedding layer as a linear layer.
to_quantized([group_size, bits])	Return a QuantizedEmbedding layer that approximates this embedding layer.