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

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

from typing import Any, Callable, Optional, Union

import mlx.core as mx
from base import Module

def quantize(
    model: Module,
    group_size: int = ...,
    bits: int = ...,
    *,
    mode: str = ...,
    class_predicate: Optional[
        Callable[[str, Module], Union[bool, dict[str, Any]]]
    ] = ...,
) -> None:
    """Quantize the sub-modules of a module according to a predicate.

    By default all layers that define a ``to_quantized(group_size, bits)``
    method will be quantized. Both :obj:`Linear` and :obj:`Embedding` layers
    will be quantized. Note also, the module is updated in-place.

    Args:
        model (Module): The model whose leaf modules may be quantized.
        group_size (int): The quantization group size (see
           :func:`mlx.core.quantize`). Default: ``64``.
        bits (int): The number of bits per parameter (see
           :func:`mlx.core.quantize`). Default: ``4``.
        mode (str): The quantization method to use (see
           :func:`mlx.core.quantize`). Default: ``"affine"``.
        class_predicate (Optional[Callable]): A callable which receives the
          :obj:`Module` path and :obj:`Module` itself and returns ``True`` or a
          dict of params for `to_quantized` if it should be quantized and
          ``False`` otherwise. If ``None``, then all layers that define a
          ``to_quantized(group_size, bits)`` method are quantized.
          Default: ``None``.
    """

class QuantizedEmbedding(Module):
    """The same as :obj:`Embedding` but with a  quantized weight matrix.

    :obj:`QuantizedEmbedding` also provides a :meth:`from_embedding`
    classmethod to convert embedding layers to :obj:`QuantizedEmbedding`
    layers.

    Args:
        num_embeddings (int): How many possible discrete tokens can we embed.
           Usually called the vocabulary size.
        dims (int): The dimensionality of the embeddings.
        group_size (int, optional): The group size to use for the quantized
            weight. See :func:`~mlx.core.quantize`. Default: ``64``.
        bits (int, optional): The bit width to use for the quantized weight.
            See :func:`~mlx.core.quantize`. Default: ``4``.
        mode (str): The quantization method to use (see
           :func:`mlx.core.quantize`). Default: ``"affine"``.
    """
    def __init__(
        self,
        num_embeddings: int,
        dims: int,
        group_size: int = ...,
        bits: int = ...,
        mode: str = ...,
    ) -> None: ...
    def __call__(self, x: mx.array) -> mx.array: ...
    def as_linear(self, x: mx.array) -> mx.array:
        """
        Call the quantized embedding layer as a quantized linear layer.

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

    @classmethod
    def from_embedding(
        cls,
        embedding_layer: Module,
        group_size: int = ...,
        bits: int = ...,
        mode: str = ...,
    ) -> QuantizedEmbedding:
        """Create a :obj:`QuantizedEmbedding` layer from an :obj:`Embedding` layer."""

class QuantizedLinear(Module):
    """Applies an affine transformation to the input using a quantized weight matrix.

    It is the quantized equivalent of :class:`Linear`. For now its
    parameters are frozen and will not be included in any gradient computation
    but this will probably change in the future.

    :obj:`QuantizedLinear` also provides a classmethod :meth:`from_linear` to
    convert linear layers to :obj:`QuantizedLinear` layers.

    Args:
        input_dims (int): The dimensionality of the input features.
        output_dims (int): The dimensionality of the output features.
        bias (bool, optional): If set to ``False`` then the layer will not use
            a bias. Default: ``True``.
        group_size (int, optional): The group size to use for the quantized
            weight. See :func:`~mlx.core.quantize`. Default: ``64``.
        bits (int, optional): The bit width to use for the quantized weight.
            See :func:`~mlx.core.quantize`. Default: ``4``.
        mode (str): The quantization method to use (see
           :func:`mlx.core.quantize`). Default: ``"affine"``.
    """
    def __init__(
        self,
        input_dims: int,
        output_dims: int,
        bias: bool = ...,
        group_size: int = ...,
        bits: int = ...,
        mode: str = ...,
    ) -> None: ...
    def __call__(self, x: mx.array) -> mx.array: ...
    @classmethod
    def from_linear(
        cls,
        linear_layer: Module,
        group_size: int = ...,
        bits: int = ...,
        mode: str = ...,
    ) -> QuantizedLinear:
        """Create a :obj:`QuantizedLinear` layer from a :obj:`Linear` layer."""