Embedding layer — mlx_embedding • Rmlx

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(). Default: mlx_default_device().

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 (0-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)