Repository navigation
Fix Glorot/He uniform init docstrings to label the uniform bound, not sigma - #3831
Conversation
Pablosinyores
left a comment
Pablosinyores
left a comment
There was a problem hiding this comment.
Confirmed correct against the implementation. Both glorot_uniform and he_uniform compute limit = gain * math.sqrt(6.0 / (fan_in + fan_out)) (resp. sqrt(3.0 / fan)) and then mx.random.uniform(-limit, limit, ...) — so that quantity is the half-width of the uniform interval, not a standard deviation. Labeling it \sigma overstates the spread by a factor of √3 (for U[-a, a] the actual std is a/√3), so the relabel is right. The explicit [-a, a] interval line you added makes the meaning unambiguous, which is the better framing.
Worth noting the normal siblings (glorot_normal / he_normal) correctly keep \sigma alongside mx.random.normal(scale=std) — so this reads like the \sigma label was carried over from the normal variants into the uniform ones by mistake.
One small thought: the implementation names this variable limit, so using \text{limit} as the symbol (as your #3835 does) ties the formula straight to the source. a with the explicit interval is also perfectly clear and arguably more conventional — either works.
Heads up that this overlaps with your own #3835 — same two functions, same lines — so they'll conflict and one probably wants closing. This is the more complete of the two (the interval prose is the useful part), so if you keep one, I'd keep this.
d0ffbbc to
769e1d8
Compare
|
Thanks for the careful review. Agreed the |
Pablosinyores
left a comment
Pablosinyores
left a comment
There was a problem hiding this comment.
Accurate fix. Both initializers compute limit = gain * sqrt(...) then mx.random.uniform(-limit, limit, ...), so labeling that bound \sigma (a std-dev symbol) was wrong for a uniform range. Renaming to \text{limit} and stating the interval [-limit, limit] matches the variable that's literally named limit, and the sqrt(6/(fan_in+fan_out)) / sqrt(3/fan) factors are the variance-match to the normal variants. LGTM.
|
Gentle ping @angeloskath. This is a two-line docstring fix, CI is green, and it picked up a community review confirming the diagnosis.
Happy to rebase if it has gone stale. |
…igma 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].
769e1d8 to
6cd5242
Compare
Proposed changes
The
glorot_uniformandhe_uniformdocstrings inpython/mlx/nn/init.pywrite the sampling bound with the symbol
\sigma, but both initializers drawfrom a symmetric uniform distribution,
uniform(-a, a), so the value shown isthe bound
a, not a standard deviation.\sigmaconventionally denotes thestandard deviation, and for
uniform(-a, a)the standard deviation isa / sqrt(3):glorot_uniform: bounda = gain * sqrt(6 / (fan_in + fan_out)), so theactual standard deviation is
gain * sqrt(2 / (fan_in + fan_out)), which isexactly the
\sigmadocumented forglorot_normal.he_uniform: bounda = gain * sqrt(3 / fan), so the actual standarddeviation is
gain / sqrt(fan), which is exactly the\sigmadocumented forhe_normal.So the documented
\sigmafor the uniform initializers is off by a factor ofsqrt(3)from the real standard deviation. It is also inconsistent with thesurrounding prose, which already describes the quantity as a "range", and with
the normal initializers in the same file, where
\sigmacorrectly denotes thestandard deviation. For comparison, PyTorch's
kaiming_uniform_docstringlabels the same quantity
boundand samplesU(-bound, bound).This relabels the quantity as the bound
aand shows the sampling interval[-a, a]. It is a documentation-only change with no effect on behavior; thesampled values are unchanged.
Checklist
pre-commit run --all-filesto format my code / installed pre-commit prior to committing changes (docstring text only; verified it compiles and does not changeblackformatting)