Prevent integer overflow in tail ESS (#133)

* Add failing test * Fix failing test * Increment patch number * Fix comment
TuringLang · Jan 8, 2025 · 5a32333 · 5a32333 · sethaxen · Jan 8, 2025
1 parent 438da79
commit 5a32333
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 2 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "MCMCDiagnosticTools"
 uuid = "be115224-59cd-429b-ad48-344e309966f0"
 authors = ["David Widmann", "Seth Axen"]
-version = "0.3.13"
+version = "0.3.14"
 
 [deps]
 AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"

diff --git a/src/ess_rhat.jl b/src/ess_rhat.jl
@@ -289,7 +289,7 @@ function _ess(
     ::Val{:tail}, x::AbstractArray{<:Union{Missing,Real}}; tail_prob::Real=1//10, kwargs...
 )
     # workaround for https://github.com/JuliaStats/Statistics.jl/issues/136
-    T = Base.promote_eltype(x, tail_prob)
+    T = float(Base.promote_eltype(x, tail_prob))
     pl = convert(T, tail_prob / 2)
     pu = convert(T, 1 - tail_prob / 2)
     S_lower = _ess(Base.Fix2(Statistics.quantile, pl), x; kwargs...)

diff --git a/test/ess_rhat.jl b/test/ess_rhat.jl
@@ -363,6 +363,17 @@ mymean(x) = mean(x)
         @test all(>(rhat_cutoff), R_tail)
     end
 
+    # check that issue #132 is fixed
+    @testset "tail ESS does not overflow for integer types" begin
+        x = rand(1:10_000, 1_000, 4, 100)
+        S, _ = ess_rhat(x; kind=:tail)
+        @test all(isfinite, S)
+        @test eltype(S) == Float64
+        S = ess(x; kind=:tail)
+        @test all(isfinite, S)
+        @test eltype(S) == Float64
+    end
+
     @testset "bulk and tail ESS and R-hat for heavy tailed" begin
         # sampling Cauchy distribution with large max depth to allow for better tail
         # exploration. From https://avehtari.github.io/rhat_ess/rhat_ess.html chains have