Embedding layer

Maps discrete tokens to continuous vectors.

Usage

mlx_embedding(num_embeddings, embedding_dim)

Arguments

num_embeddings

Size of vocabulary.

embedding_dim

Dimension of embedding vectors.

Value

An mlx_module for token embeddings.

See also

mlx.nn.Embedding

Examples

set.seed(1)
emb <- mlx_embedding(num_embeddings = 100, embedding_dim = 16)
# Token indices (1-indexed)
tokens <- as_mlx(matrix(c(5, 10, 3, 7), 2, 2))
mlx_forward(emb, tokens)
#> mlx array [2 x 2 x 16]
#>   dtype: float32
#>   (64 elements, not shown)