Embedding layer

Maps discrete tokens to continuous vectors.

Usage

mlx_embedding(num_embeddings, embedding_dim, device = mlx_default_device())

Arguments

num_embeddings

Size of vocabulary.

embedding_dim

Dimension of embedding vectors.

device

Execution target: supply "gpu", "cpu", or an mlx_stream created via mlx_new_stream(). Defaults to the current mlx_default_device() unless noted otherwise (helpers that act on an existing array typically reuse that array's device or stream).

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
#>   device: gpu
#>   (64 elements, not shown)