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

from typing import Callable, Literal

import mlx.core as mx

def constant(value: float, dtype: mx.Dtype = ...) -> Callable[[mx.array], mx.array]:
    r"""An initializer that returns an array filled with ``value``.

    Args:
        value (float): The value to fill the array with.
        dtype (Dtype, optional): The data type of the array. Default:
          ``float32``.

    Returns:
        Callable[[array], array]: An initializer that returns an array with the
        same shape as the input, filled with ``value``.

    Example:

        >>> init_fn = nn.init.constant(0.5)
        >>> init_fn(mx.zeros((2, 2)))
        array([[0.5, 0.5],
               [0.5, 0.5]], dtype=float32)
    """

def normal(
    mean: float = ..., std: float = ..., dtype: mx.Dtype = ...
) -> Callable[[mx.array], mx.array]:
    r"""An initializer that returns samples from a normal distribution.

    Args:
        mean (float, optional): Mean of the normal distribution. Default:
          ``0.0``.
        std (float, optional): Standard deviation of the normal distribution.
          Default: ``1.0``.
        dtype (Dtype, optional): The data type of the array. Default:
          ``float32``.

    Returns:
        Callable[[array], array]: An initializer that returns an array with the
        same shape as the input, filled with samples from a normal distribution.

    Example:

        >>> init_fn = nn.init.normal()
        >>> init_fn(mx.zeros((2, 2)))
        array([[-0.982273, -0.534422],
               [0.380709, 0.0645099]], dtype=float32)
    """

def uniform(
    low: float = ..., high: float = ..., dtype: mx.Dtype = ...
) -> Callable[[mx.array], mx.array]:
    r"""An initializer that returns samples from a uniform distribution.

    Args:
        low (float, optional): The lower bound of the uniform distribution.
          Default: ``0.0``.
        high (float, optional): The upper bound of the uniform distribution.
          Default: ``1.0``
        dtype (Dtype, optional): The data type of the array. Default: ``float32``.

    Returns:
        Callable[[array], array]: An initializer that returns an array
        with the same shape as the input, filled with samples from a uniform
        distribution

    Example:

        >>> init_fn = nn.init.uniform(low=0, high=1)
        >>> init_fn(mx.zeros((2, 2)))
        array([[0.883935, 0.863726],
               [0.617261, 0.417497]], dtype=float32)
    """

def identity(dtype: mx.Dtype = ...) -> Callable[[mx.array], mx.array]:
    r"""An initializer that returns an identity matrix.

    Args:
        dtype (Dtype, optional): The data type of the array. Defaults:
          ``float32``.

    Returns:
        Callable[[array], array]: An initializer that returns an identity
        matrix with the same shape as the input.

    Example:

        >>> init_fn = nn.init.identity()
        >>> init_fn(mx.zeros((2, 2)))
        array([[1, 0],
               [0, 1]], dtype=float32)
    """

def glorot_normal(dtype: mx.Dtype = ...) -> Callable[[mx.array, float], mx.array]:
    r"""A Glorot normal initializer.

    This initializer samples from a normal distribution with a standard
    deviation computed from the number of input (``fan_in``) and output
    (``fan_out``) units according to:

    .. math::
        \sigma = \gamma \sqrt{\frac{2.0}{\text{fan\_in} + \text{fan\_out}}}

    For more details see the original reference: `Understanding the difficulty
    of training deep feedforward neural networks
    <https://proceedings.mlr.press/v9/glorot10a.html>`_

    Args:
        dtype (Dtype, optional): The data type of the array. Default: ``float32``.

    Returns:
        Callable[[array, float], array]: An initializer that returns an array
        with the same shape as the input, filled with samples from the Glorot
        normal distribution.

    Example:

        >>> init_fn = nn.init.glorot_normal()
        >>> init_fn(mx.zeros((2, 2)))
        array([[0.191107, 1.61278],
               [-0.150594, -0.363207]], dtype=float32)
        >>> init_fn(mx.zeros((2, 2)), gain=4.0)
        array([[1.89613, -4.53947],
               [4.48095, 0.995016]], dtype=float32)
    """

def glorot_uniform(dtype: mx.Dtype = ...) -> Callable[[mx.array, float], mx.array]:
    r"""A Glorot uniform initializer.

    This initializer samples from a uniform distribution with a range
    computed from the number of input (``fan_in``) and output (``fan_out``)
    units according to:

    .. math::
        \sigma = \gamma \sqrt{\frac{6.0}{\text{fan\_in} + \text{fan\_out}}}

    For more details see the original reference: `Understanding the difficulty
    of training deep feedforward neural networks
    <https://proceedings.mlr.press/v9/glorot10a.html>`_

    Args:
        dtype (Dtype, optional): The data type of the array. Default: ``float32``.

    Returns:
        Callable[[array, float], array]: An initializer that returns an array
        with the same shape as the input, filled with samples from the Glorot
        uniform distribution.

    Example:

        >>> init_fn = nn.init.glorot_uniform()
        >>> init_fn(mx.zeros((2, 2)))
        array([[0.223404, -0.890597],
               [-0.379159, -0.776856]], dtype=float32)
        >>> init_fn(mx.zeros((2, 2)), gain=4.0)
        array([[-1.90041, 3.02264],
               [-0.912766, 4.12451]], dtype=float32)
    """

def he_normal(
    dtype: mx.Dtype = ...,
) -> Callable[[mx.array, Literal["fan_in", "fan_out"], float], mx.array]:
    r"""Build a He normal initializer.

    This initializer samples from a normal distribution with a standard
    deviation computed from the number of input (``fan_in``) or output
    (``fan_out``) units according to:

    .. math::
        \sigma = \gamma \frac{1}{\sqrt{\text{fan}}}

    where :math:`\text{fan}` is either the number of input units when the
    ``mode`` is ``"fan_in"`` or output units when the ``mode`` is
    ``"fan_out"``.

    For more details see the original reference: `Delving Deep into Rectifiers:
    Surpassing Human-Level Performance on ImageNet Classification
    <https://arxiv.org/abs/1502.01852>`_

    Args:
        dtype (Dtype, optional): The data type of the array. Defaults to mx.float32.

    Returns:
        Callable[[array, str, float], array]: An initializer that returns an
        array with the same shape as the input, filled with samples from the He
        normal distribution.

    Example:

        >>> init_fn = nn.init.he_normal()
        >>> init_fn(mx.zeros((2, 2)))  # uses fan_in
        array([[-1.25211, 0.458835],
               [-0.177208, -0.0137595]], dtype=float32)
        >>> init_fn(mx.zeros((2, 2)), mode="fan_out", gain=5)
        array([[5.6967, 4.02765],
               [-4.15268, -2.75787]], dtype=float32)
    """

def he_uniform(
    dtype: mx.Dtype = ...,
) -> Callable[[mx.array, Literal["fan_in", "fan_out"], float], mx.array]:
    r"""A He uniform (Kaiming uniform) initializer.

    This initializer samples from a uniform distribution with a range
    computed from the number of input (``fan_in``) or output (``fan_out``)
    units according to:

    .. math::

        \sigma = \gamma \sqrt{\frac{3.0}{\text{fan}}}

    where :math:`\text{fan}` is either the number of input units when the
    ``mode`` is ``"fan_in"`` or output units when the ``mode`` is
    ``"fan_out"``.

    For more details see the original reference: `Delving Deep into Rectifiers:
    Surpassing Human-Level Performance on ImageNet Classification
    <https://arxiv.org/abs/1502.01852>`_


    Args:
        dtype (Dtype, optional): The data type of the array. Default: ``float32``.

    Returns:
        Callable[[array, str, float], array]: An initializer that returns an
        array with the same shape as the input, filled with samples from  the
        He uniform distribution.

    Example:

        >>> init_fn = nn.init.he_uniform()
        >>> init_fn(mx.zeros((2, 2)))  # uses fan_in
        array([[0.0300242, -0.0184009],
               [0.793615, 0.666329]], dtype=float32)
        >>> init_fn(mx.zeros((2, 2)), mode="fan_out", gain=5)
        array([[-1.64331, -2.16506],
               [1.08619, 5.79854]], dtype=float32)
    """

def sparse(
    sparsity: float, mean: float = ..., std: float = ..., dtype: mx.Dtype = ...
) -> Callable[[mx.array], mx.array]:
    r"""An initializer that returns a sparse matrix.

    Args:
        sparsity (float): The fraction of elements in each column to be set to
        zero.
        mean (float, optional): Mean of the normal distribution. Default:
          ``0.0``.
        std (float, optional): Standard deviation of the normal distribution.
          Default: ``1.0``.
        dtype (Dtype, optional): The data type of the array. Default:
          ``float32``.

    Returns:
        Callable[[array], array]: An initializer that returns an array with the
        same shape as the input, filled with samples from a normal distribution.

    Example:

        >>> init_fn = nn.init.sparse(sparsity=0.5)
        >>> init_fn(mx.zeros((2, 2)))
        array([[-1.91187, -0.117483],
       [0, 0]], dtype=float32)
    """

def orthogonal(
    gain: float = ..., dtype: mx.Dtype = ...
) -> Callable[[mx.array], mx.array]:
    r"""An initializer that returns an orthogonal matrix.

    Args:
        gain (float, optional): Scaling factor for the orthogonal matrix.
            Default: ``1.0``.
        dtype (Dtype, optional): Data type of the array. Default: ``float32``.

    Returns:
        Callable[[array], array]: An initializer that returns
        an orthogonal matrix with the same shape as the input.
    """
