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

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

from typing import Any

import mlx.core as mx
from base import Module

from .quantized import QuantizedLinear

class Identity(Module):
    r"""A placeholder identity operator that is argument-insensitive.

    Args:
        args: any argument (unused)
        kwargs: any keyword argument (unused)
    """
    def __init__(self, *args: Any, **kwargs: Any) -> None: ...
    def __call__(self, x: mx.array) -> mx.array: ...

class Linear(Module):
    r"""Applies an affine transformation to the input.

    Concretely:

    .. math::

        y = x W^\top + b

    where:
    where :math:`W` has shape ``[output_dims, input_dims]`` and :math:`b` has shape ``[output_dims]``.

    The values are initialized from the uniform distribution :math:`\mathcal{U}(-{k}, {k})`,
    where :math:`k = \frac{1}{\sqrt{D_i}}` and :math:`D_i` is equal to ``input_dims``.

    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 is ``True``.
    """

    weight: mx.array
    bias: mx.array | None

    def __init__(self, input_dims: int, output_dims: int, bias: bool = ...) -> None: ...
    def __call__(self, x: mx.array) -> mx.array: ...
    def to_quantized(
        self, group_size: int = ..., bits: int = ..., mode: str = ...
    ) -> QuantizedLinear:
        """Return a :obj:`QuantizedLinear` layer that approximates this layer."""

class Bilinear(Module):
    r"""Applies a bilinear transformation to the inputs.

    Concretely:

    .. math::

        y_i = x_1^\top W_i x_2 + b_i

    where:
    :math:`W` has shape ``[output_dims, input1_dims, input2_dims]``, :math:`b` has shape ``[output_dims ]``,
    and :math:`i` indexes the output dimension.

    The values are initialized from the uniform distribution :math:`\mathcal{U}(-{k}, {k})`,
    where :math:`k = \frac{1}{\sqrt{D_1}}` and :math:`D_1` is ``input1_dims``.

    Args:
        input1_dims (int): The dimensionality of the input1 features
        input2_dims (int): The dimensionality of the input2 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 is ``True``.
    """
    def __init__(
        self, input1_dims: int, input2_dims: int, output_dims: int, bias: bool = ...
    ) -> None: ...
    def __call__(self, x1: mx.array, x2: mx.array) -> mx.array: ...