Reduce allocations in all + any and remove isfinite type piracy (#388)

* Reduce allocations in `all` + `any` and remove `isfinite` type piracy

* Fix `isfinite` check

Author: devmotion

research/src/riemannian_hmc.jl

@@ -356,7 +356,7 @@ function ∂H∂r(
     r::AbstractVecOrMat,
 )
     H = h.metric.G(θ)
-    # if any(.!(isfinite.(H)))
+    # if !all(isfinite, H)
     #     println("θ: ", θ)
     #     println("H: ", H)
     # end

research/src/riemannian_hmc_utility.jl

@@ -43,7 +43,7 @@ function prepare_sample_target(hps, θ₀, ℓπ)
     fstabilize = H -> H + hps.λ * I
     Gfunc = x -> begin
         H = fstabilize(Hfunc(x)[3])
-        any(.!(isfinite.(H))) ? diagm(ones(length(x))) : H
+        all(isfinite, H) ? H : diagm(ones(length(x)))
     end
     _∂G∂θfunc = gen_∂G∂θ_fwd(Vfunc, θ₀; f = fstabilize) # size==(4, 2)
     ∂G∂θfunc = x -> reshape_∂G∂θ(_∂G∂θfunc(x)) # size==(2, 2, 2)

research/tests/riemannian_hmc.jl

@@ -43,7 +43,7 @@ using AdvancedHMC: neg_energy, energy
             hamiltonian = Hamiltonian(metric, kinetic, ℓπ, ∂ℓπ∂θ)
 
             if hessmap isa SoftAbsMap || # only test kinetic energy for SoftAbsMap as that of IdentityMap can be non-PD
-               all(iszero.(x)) # or for x==0 that I know it's PD
+               all(iszero, x) # or for x==0 that I know it's PD
                 @testset "Kinetic energy" begin
                     Σ = hamiltonian.metric.map(hamiltonian.metric.G(x))
                     @test neg_energy(hamiltonian, r, x) ≈ logpdf(MvNormal(zeros(D), Σ), r)

src/adaptation/stepsize.jl

@@ -127,7 +127,7 @@ function adapt_stepsize!(
     @debug "Adapting step size..." new_ϵ = ϵ old_ϵ = da.state.ϵ
 
     # TODO: we might want to remove this when all other numerical issues are correctly handelled
-    if any(isnan.(ϵ)) || any(isinf.(ϵ))
+    if !all(isfinite, ϵ)
         @warn "Incorrect ϵ = $ϵ; ϵ_previous = $(da.state.ϵ) is used instead."
         # FIXME: this revert is buggy for batch mode
         @unpack m, ϵ, x_bar, H_bar = state

src/hamiltonian.jl

@@ -58,15 +58,15 @@ struct PhasePoint{T<:AbstractVecOrMat{<:AbstractFloat},V<:DualValue}
     ℓκ::V # Cached neg kinect energy for the current r.
     function PhasePoint(θ::T, r::T, ℓπ::V, ℓκ::V) where {T,V}
         @argcheck length(θ) == length(r) == length(ℓπ.gradient) == length(ℓπ.gradient)
-        if any(isfinite.((θ, r, ℓπ, ℓκ)) .== false)
-            # @warn "The current proposal will be rejected due to numerical error(s)." isfinite.((θ, r, ℓπ, ℓκ))
-            # NOTE eltype has to be inlined to avoid type stability issue; see #267
+        if !isfinite(ℓπ)
             ℓπ = DualValue(
-                map(v -> isfinite(v) ? v : -eltype(T)(Inf), ℓπ.value),
+                map(v -> isfinite(v) ? v : oftype(v, -Inf), ℓπ.value),
                 ℓπ.gradient,
             )
+        end
+        if !isfinite(ℓκ)
             ℓκ = DualValue(
-                map(v -> isfinite(v) ? v : -eltype(T)(Inf), ℓκ.value),
+                map(v -> isfinite(v) ? v : oftype(v, -Inf), ℓκ.value),
                 ℓκ.gradient,
             )
         end
@@ -105,7 +105,6 @@ end
 
 
 Base.isfinite(v::DualValue) = all(isfinite, v.value) && all(isfinite, v.gradient)
-Base.isfinite(v::AbstractVecOrMat) = all(isfinite, v)
 Base.isfinite(z::PhasePoint) = isfinite(z.ℓπ) && isfinite(z.ℓκ)
 
 ###

src/trajectory.jl

@@ -275,7 +275,7 @@ function transition(
             hamiltonian_energy = H,
             hamiltonian_energy_error = H - H0,
             # check numerical error in proposed phase point. 
-            numerical_error = isfinite(H′),
+            numerical_error = !all(isfinite, H′),
         ),
         stat(τ.integrator),
     )
@@ -296,13 +296,12 @@ function accept_phasepoint!(
 end
 function accept_phasepoint!(z::T, z′::T, is_accept) where {T<:PhasePoint{<:AbstractMatrix}}
     # Revert unaccepted proposals in `z′`
-    is_reject = (!).(is_accept)
-    if any(is_reject)
+    if !all(is_accept)
         # Convert logical indexing to number indexing to support CUDA.jl
         # NOTE: for x::CuArray, x[:,Vector{Bool}]  is NOT supported
         #                       x[:,CuVector{Int}] is NOT supported
         #                       x[:,Vector{Int}]   is     supported
-        is_reject = findall(is_reject) |> Array
+        is_reject = Vector(findall(!, is_accept))
         z′.θ[:, is_reject] = z.θ[:, is_reject]
         z′.r[:, is_reject] = z.r[:, is_reject]
         z′.ℓπ.value[is_reject] = z.ℓπ.value[is_reject]

test/sampler.jl

@@ -23,7 +23,7 @@ function test_stats(
         :hamiltonian_energy_error,
         :is_adapt,
     )
-        @test all(map(s -> in(name, propertynames(s)), stats))
+        @test all(s -> in(name, propertynames(s)), stats)
     end
     is_adapts = getproperty.(stats, :is_adapt)
     @test is_adapts[1:n_adapts] == ones(Bool, n_adapts)
@@ -49,7 +49,7 @@ function test_stats(
         :tree_depth,
         :numerical_error,
     )
-        @test all(map(s -> in(name, propertynames(s)), stats))
+        @test all(s -> in(name, propertynames(s)), stats)
     end
     is_adapts = getproperty.(stats, :is_adapt)
     @test is_adapts[1:n_adapts] == ones(Bool, n_adapts)