← back to Exo
.typings/mlx/nn/init.pyi
285 lines
"""
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.
"""