mlx.nn.Embedding

class mlx.nn.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.