Work around compiler bug - Julia issue 34232

src/trajectory.jl

@@ -641,21 +641,29 @@ find_good_eps(
     max_n_iters::Int=100
 ) where {T<:AbstractFloat} = find_good_eps(GLOBAL_RNG, h, θ; max_n_iters=max_n_iters)
 
-_rand(rng::AbstractRNG) = rand(rng)
-_rand(rng::AbstractVector{<:AbstractRNG}) = rand.(rng)
-
 """
 Perform MH acceptance based on energy, i.e. negative log probability.
 """
 function mh_accept_ratio(
-    rng::Union{AbstractRNG,AbstractVector{<:AbstractRNG}},
-    Horiginal::AbstractScalarOrVec{<:T},
-    Hproposal::AbstractScalarOrVec{<:T}
+    rng::AbstractRNG,
+    Horiginal::T,
+    Hproposal::T,
+) where {T <: AbstractFloat}
+    α = min(1.0, exp(Horiginal - Hproposal))
+    accept = rand(rng) < α
+    return accept, α
+end
+function mh_accept_ratio(
+    rng::Union{AbstractRNG, AbstractVector{<:AbstractRNG}},
+    Horiginal::AbstractVector{<:T},
+    Hproposal::AbstractVector{<:T},
 ) where {T<:AbstractFloat}
     α = min.(1.0, exp.(Horiginal .- Hproposal))
     accept = _rand(rng) .< α
     return accept, α
 end
+_rand(rng::AbstractRNG) = rand(rng)
+_rand(rng::AbstractVector{<:AbstractRNG}) = rand.(rng)
 
 ####
 #### Adaption

test/adaptation/precond.jl

@@ -101,8 +101,8 @@ let D=10
         return ℓπ, ∂ℓπ∂θ
     end
 
-    function runnuts(ℓπ, ∂ℓπ∂θ, metric; n_samples=2_000)
-        n_adapts = 1_000
+    function runnuts(ℓπ, ∂ℓπ∂θ, metric; n_samples=3_000)
+        n_adapts = 1_500
 
         θ_init = rand(D)