Improve type stability for truncated normal moments #1717
Conversation
The internal functions used for computing the mean and variance of the
truncated normal distribution were implemented using literals like `√2`,
which are always `Float64`. However, some branches of the functions
based the return type off of the types of the arguments, which needn't
be `Float64`. This led to some instability that can be seen via
`@code_warntype`. For example, the inferred result of `_tnmom1` for
`Float32` inputs was `Union{Float32,Float64}`.
We can instead use the irrational constants defined for square roots of
pi and 2 and ratios thereof in order to avoid promoting to `Float64`
unnecessarily. This makes the return type concretely inferrable and
provides a small performance improvement, about 7 ns for `mean` and 20
ns for `var` on my machine as compared to current master.
Codecov ReportPatch coverage:
Additional details and impacted files@@ Coverage Diff @@
## master #1717 +/- ##
=======================================
Coverage 85.82% 85.83%
=======================================
Files 137 137
Lines 8318 8322 +4
=======================================
+ Hits 7139 7143 +4
Misses 1179 1179
☔ View full report in Codecov by Sentry. |
Co-authored-by: David Widmann <devmotion@users.noreply.github.com>
| end | ||
| Δ = (b - a) * middle(a, b) | ||
| Δ = (b - a) * mid | ||
| a′ = a * invsqrt2 |
There was a problem hiding this comment.
@devmotion, I'm curious if you have any thoughts on this approach vs. a / sqrt2. They seem to be extremely close but not exactly identical (but close enough to not affect the tests). Similar situation for the division of things by sqrthalfπ vs. defining a separate irrational for sqrt(2/π) and multiplying.
There was a problem hiding this comment.
Hm no, to me both approaches seem fine. I guess if an irrational of the inverse exists, multiplication might be faster and therefore preferable.
There was a problem hiding this comment.
That was my thought as well, just wasn't sure whether there was something I might be missing in terms of potential loss of precision or things like that
The internal functions used for computing the mean and variance of the truncated normal distribution were implemented using literals like
√2, which are alwaysFloat64. However, some branches of the functions based the return type off of the types of the arguments, which needn't beFloat64. This led to some instability that can be seen via@code_warntype. For example, the inferred result of_tnmom1forFloat32inputs wasUnion{Float32,Float64}.We can instead use the irrational constants defined for square roots of pi and 2 and ratios thereof in order to avoid promoting to
Float64unnecessarily. This makes the return type concretely inferrable and provides a small performance improvement, about 12% better (-7ns) formeanand 14% better (-20ns) forvaron my machine as compared to current master.