From 769e1d87a39be1cff56674b5e87c30570e35edad Mon Sep 17 00:00:00 2001 From: Vineeth Sai Date: Fri, 10 Jul 2026 13:38:13 -0700 Subject: [PATCH] Fix Glorot/He uniform init docstrings: label the uniform bound, not sigma The glorot_uniform and he_uniform docstrings wrote the sampling bound as sigma (the standard-deviation symbol), but the code samples from uniform(-a, a) where a is that bound. The actual standard deviation of the result is a / sqrt(3), so labeling the bound sigma is off by sqrt(3) and is inconsistent with the prose (which calls it a range) and with the normal initializers, where sigma correctly denotes the standard deviation. Relabel the quantity as the bound a and show the sampling interval [-a, a]. --- python/mlx/nn/init.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/python/mlx/nn/init.py b/python/mlx/nn/init.py index 3fdf42990b..f0bda14a66 100644 --- a/python/mlx/nn/init.py +++ b/python/mlx/nn/init.py @@ -194,12 +194,13 @@ def glorot_uniform( ) -> 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``) + This initializer samples from a uniform distribution on the interval + :math:`[-\text{limit}, \text{limit}]`, where the bound :math:`\text{limit}` + is 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}}} + \text{limit} = \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 @@ -295,13 +296,14 @@ def he_uniform( ) -> 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``) + This initializer samples from a uniform distribution on the interval + :math:`[-\text{limit}, \text{limit}]`, where the bound :math:`\text{limit}` + is computed from the number of input (``fan_in``) or output (``fan_out``) units according to: .. math:: - \sigma = \gamma \sqrt{\frac{3.0}{\text{fan}}} + \text{limit} = \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