`.typings/mlx/nn/layers/embedding.pyi`

35 lines
"""
This type stub file was generated by pyright.
"""

import mlx.core as mx
from base import Module

from .quantized import QuantizedEmbedding

class Embedding(Module):
    """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.

    Args:
        num_embeddings (int): How many possible discrete tokens can we embed.
           Usually called the vocabulary size.
        dims (int): The dimensionality of the embeddings.
    """
    def __init__(self, num_embeddings: int, dims: int) -> None: ...
    def __call__(self, x: mx.array) -> mx.array: ...
    def as_linear(self, x: mx.array) -> mx.array:
        """
        Call the embedding layer as a linear layer.

        Use this for example when input embedding and output projection
        weights are tied.
        """

    def to_quantized(
        self, group_size: int = ..., bits: int = ..., mode: str = ...
    ) -> QuantizedEmbedding:
        """Return a :obj:`QuantizedEmbedding` layer that approximates this embedding layer."""