diff --git a/src/HVIApproximation.jl b/src/HVIApproximation.jl index 8a78539..d3a46b5 100644 --- a/src/HVIApproximation.jl +++ b/src/HVIApproximation.jl @@ -34,6 +34,8 @@ struct MeanHVIApproximationDev <: AbstractMeanHVIApproximation end # neglecting correlations struct DiagonalHVIApproximation <: AbstractHVIApproximation end +get_numberof_θM(::DiagonalHVIApproximation, ml_pred::AbstractArray) = length(ml_pred) + abstract type AbstractMeanVarSepHVIApproximation <: AbstractHVIApproximation end diff --git a/src/elbo_site.jl b/src/elbo_site.jl index 8c1846f..8a21430 100644 --- a/src/elbo_site.jl +++ b/src/elbo_site.jl @@ -6,6 +6,13 @@ function randnPM!(rng, rnorm::NamedTuple) nothing end +""" + Dispatch to make some part of neg_elbo_sites! polymorphic +""" +abstract type AbstractElboLossMode end +struct NegElbo <: AbstractElboLossMode end +struct NegLogDensity <: AbstractElboLossMode end + """ elbo_helpers need to be initialized with new random numbers in h.ζsP and hi.ζsM by calling randnPM! before. @@ -13,6 +20,7 @@ By this way, we can compute the derivative corresponding to the forward pass """ function neg_elbo_sites!( elbo_helpers::NamedTuple, # tuple of preallocated arrays + loss_mode::AbstractElboLossMode, ϕ::NamedTuple, rnormPM::NamedTuple, sample_args::NamedTuple, @@ -34,7 +42,7 @@ function neg_elbo_sites!( ϕqIc = intϕqI(ϕqI) pbm_covar_indices = !isnothing(pbm_covar_indices) && isempty(pbm_covar_indices) ? nothing : pbm_covar_indices # - sample_ζsP!(h.ζsP, h.logσ_ζP, approx, rnormPM.P, ϕqPc, cor_ends.P, h.sample_buffers) # n_P * n_MC + set_ζsP!(h.ζsP, h.logσ_ζP, loss_mode, approx, rnormPM.P, ϕqPc, cor_ends.P, h.sample_buffers) # n_P * n_MC # (n_M x n_sit) or (n_M x n_MC x n_sit) ϕms_buffer_key = isnothing(pbm_covar_indices) ? :ϕms : :ϕms_mcs #g_apply!(h[ϕms_buffer_key], ϕg, xM, h.ζsP, pbm_covar_indices, g, h.xMP, is_testmode, h.ϕms_mcs2D_buffer) @@ -49,10 +57,10 @@ function neg_elbo_sites!( function compute_nelboi_z_cl!(hi, rnormM, nljoint_args_indi_tup, ϕm) #nljoint_args_indi = NamedTuple{keys_nljoint_args_indi}(nljoint_args_indi_tup) nljoint_args_indi = name_tuple(nljoint_args_indi_tup, keys_nljoint_args_indi_val) - compute_nelboi_z!(hi, approx, rnormM, nljoint_args_indi, ϕm, + compute_nelboi_z!(hi, loss_mode, approx, rnormM, nljoint_args_indi, ϕm, ϕqIc, θsP, cor_ends.M, transM, nljoint_args_fix) end - #res_site = map(compute_nelboi_z!, h.helpers_sites, rnormPM.M, i_sites_train, ϕm_it) + #res_site = map(compute_nelboi_z!, loss_mode, h.helpers_sites, rnormPM.M, i_sites_train, ϕm_it) #MAYBE: distributed mapreduce: # https://docs.julialang.org/en/v1/stdlib/Distributed/#Distributed.@distributed # https://github.com/SupaeroDataScience/DE/blob/main/notebooks/Introduction%20to%20MapReduce.ipynb @@ -64,12 +72,26 @@ function neg_elbo_sites!( elbo = elbo_z - exp_ladJacTP - sum(h.logσ_ζP) end -name_tuple(tup, ::Val{KEYS}) where KEYS = NamedTuple{KEYS}(tup) +@inline function set_ζsP!( + ζsP, logσ_ζP, ::NegElbo, approx, rnormP, ϕqPc, cor_endsP, sample_buffers) + sample_ζsP!(ζsP, logσ_ζP, approx, rnormP, ϕqPc, cor_endsP, sample_buffers) # n_P * n_MC +end +@inline function set_ζsP!( + ζsP, logσ_ζP::AbstractMatrix{T}, ::NegLogDensity, approx, rnormP, ϕqPc, cor_endsP, sample_buffers) where T + ζsP .= CA.getdata(view(ϕqPc,Val(:μζP))) + logσ_ζP .= zero(T) # omit loss component of entropy of distribution +end +@inline function set_ζsP!( + ζsP, logσ_ζP::AbstractVector{T}, ::NegLogDensity, approx, rnormP, ϕqPc, cor_endsP, sample_buffers) where T + ζsP .= CA.getdata(view(ϕqPc,Val(:μζP))) + logσ_ζP .= zero(T) # omit loss component of entropy of distribution +end +name_tuple(tup, ::Val{KEYS}) where KEYS = NamedTuple{KEYS}(tup) -function compute_nelboi_z!(hi, approx::AbstractHVIApproximation, +function compute_nelboi_z!(hi, loss_mode::AbstractElboLossMode, approx::AbstractHVIApproximation, rnormM, nljoint_args_indi, ϕm::AbstractArray, ϕqIc::AbstractArray, θsP, cor_endsM, transM, nljoint_args_fix, ) # DM is 1 (vector) if no population covariates - constant across MC @@ -81,18 +103,34 @@ function compute_nelboi_z!(hi, approx::AbstractHVIApproximation, hi = map_leaves_nt(x -> PAT.get_tmp(x, template), hi) end #ζsM, logσ_ζM, rnorm, ϕqc::AbstractVector{T}, ϕm::AbstractMatrix, buffer_nθM::AbstractVector - sample_ζsM!(hi.ζsM, hi.logσ_ζM, approx, rnormM, ϕqIc, ϕm, cor_endsM, hi.sample_buffers) + set_ζsM!(hi.ζsM, hi.logσ_ζM, loss_mode, approx, rnormM, ϕqIc, ϕm, cor_endsM, hi.sample_buffers) exp_ladJacTM = transformζ!(hi.θsM, transM, hi.ζsM) # return value captures ladJacT # first component needs to be the full elbo exp_nL = exp_nLi!(hi.y, θsP, hi.θsM, nljoint_args_indi, nljoint_args_fix, hi)[1] nelboi_z = exp_nL - exp_ladJacTM - sum(hi.logσ_ζM) end + +@inline function set_ζsM!(ζsM, logσ_ζM, ::NegElbo, approx, rnormM, ϕqIc, ϕm, cor_endsM, sample_buffers) + sample_ζsM!(ζsM, logσ_ζM, approx, rnormM, ϕqIc, ϕm, cor_endsM, sample_buffers) +end +@inline function set_ζsM!(ζsM, logσ_ζM::AbstractArray{T}, ::NegLogDensity, approx, rnormM, ϕqIc, ϕm::AbstractVector, cor_endsM, sample_buffers) where T + ζsM .= view(ϕm, 1:size(ζsM,1)) # recycling fills all rows + logσ_ζM .= zero(T) + nothing +end +@inline function set_ζsM!(ζsM, logσ_ζM::AbstractArray{T}, ::NegLogDensity, approx, rnormM, ϕqIc, ϕm::AbstractMatrix, cor_endsM, sample_buffers) where T + ζsM .= view(ϕm, 1:size(ζsM,1), :) # μζM in first n_θM entries + logσ_ζM .= zero(T) + nothing +end + # get_tmp_rec_(x::PAT.DiffCache{<:AbstractArray}, template) = PAT.get_tmp(x, template) # function get_tmp_rec_(x::Union{Tuple,NamedTuple}, template) # map(xi -> get_tmp_rec_(xi, template), x) # end + function prepare_rnorm(::AbstractVector{TF}; n_θP, n_θM, n_site, n_MC) where TF (; P = Matrix{TF}(undef, n_θP, n_MC), diff --git a/src/elbo_site_grad.jl b/src/elbo_site_grad.jl index 7404354..c227ece 100644 --- a/src/elbo_site_grad.jl +++ b/src/elbo_site_grad.jl @@ -2,12 +2,13 @@ function grad_neg_elbo_sites( elbo_helpers::NamedTuple, # tuple of preallocated arrays grad_elbo_helpers::NamedTuple, # tuple of preallocated arrays pullback closures + loss_mode, ϕ::NamedTuple, rnormPM::NamedTuple, sample_args::NamedTuple, indiv_args::NamedTuple, nljoint_args_fix::NamedTuple; - n_MC_val::Val{NMC}, # need to pass by value for type inference on chunks + n_MC_val::Val{NMC_ELBO}, # need to pass by value for type inference on chunks # approx::AbstractHVIApproximation, # rnormPM::NamedTuple, # tuple of random numbers # ϕg::AbstractVector{TG}, ϕqP::AbstractVector{TF}, ϕqI::AbstractVector{TF}, g, @@ -21,7 +22,13 @@ function grad_neg_elbo_sites( # is_testmode, n_workers = 1, # TODO compute from executor and Distributed.nworkers and nthreads,... executor::Transducers.Executor = Transducers.SequentialEx(), -) where NMC +) where NMC_ELBO + if loss_mode isa NegLogDensity + NMC = 1 + n_MC_val = Val(1) + else + NMC = NMC_ELBO + end h = elbo_helpers # preallocated μζP, dμζP, ζsP, ϕms, xMP, dxMP (;ϕg, ϕqP, ϕqI) = ϕ TF = eltype(ϕqP) @@ -48,10 +55,11 @@ function grad_neg_elbo_sites( ϕqIc, selectdim(h_ϕm, ndims(h_ϕm), 1), h.ζsP, ϕg, ϕqPc, approx; pbm_covar_indices, n_workers, h, rnormMi = rnormPM.M[1], nljoint_args_indi, - n_site, n_cov, cor_ends, transP, transM, nljoint_args_fix, n_MC_val) + n_site, n_cov, cor_ends, transP, transM, nljoint_args_fix, loss_mode, + n_MC_val) hw_channel = gradh.hw_channel # - sample_ζsP!(h.ζsP, h.logσ_ζP, approx, rnormPM.P, ϕqPc, cor_ends.P, h.sample_buffers) # n_P * n_MC + set_ζsP!(h.ζsP, h.logσ_ζP, loss_mode, approx, rnormPM.P, ϕqPc, cor_ends.P, h.sample_buffers) # n_P * n_MC g_apply!(h_ϕm, ϕg, xM, h.ζsP, pbm_covar_indices, g, h.xMP, is_testmode) ladJacTP = transformζ!(h.θsP, transP, h.ζsP) # return value captures ladJacT # @@ -67,7 +75,7 @@ function grad_neg_elbo_sites( end cl = ForwardDiffGradNelboiZCl( approx, ϕqIc, gradh.dϕmvecs, gradh.hw_channel, cor_ends.M, transM, - Val(keys(nljoint_args_inds)), nljoint_args_fix, n_MC_val) + Val(keys(nljoint_args_inds)), nljoint_args_fix, loss_mode, n_MC_val) ϕm_it = eachslice(h_ϕm; dims = ndims(h_ϕm)) # _a1 = Test.@inferred first(zip(h.helpers_sites, rnormPM.M, zip_eachlastdims(nljoint_args_inds), ϕm_it, 1:n_site)) # Test.@inferred cl(_a1) @@ -87,7 +95,7 @@ function grad_neg_elbo_sites( copyto!(∂elbo_∂θP, res.dθsP) copyto!(gradh.dϕqIc, res.dϕqI) # other gradients are one, calculated from formula - gradh.∂elbo_∂logσ_ζP .= -one(TF) + set_∂elbo_∂logσ_ζP!(gradh.∂elbo_∂logσ_ζP, loss_mode) ∂elbo_∂ladJacTP = -one(TF) # # pullback gradients of ϕqm -> gradh.dϕg and gradh.dζsP @@ -107,18 +115,39 @@ function grad_neg_elbo_sites( ) # # pullback gradients of ∂elbo_∂ζP, ∂elbo_∂ϕm_∂ζP, and ∂elbo_∂logσ_ζP to dϕqP - gradh.pullback_cl_sample_ζsP!( - gradh.dϕqPc, + pullack_sample_ζsP!(gradh.dϕqPc, loss_mode, ∂elbo_∂θP_∂ζP + gradh.∂elbo_∂ϕm_∂ζP, gradh.∂elbo_∂logσ_ζP, + gradh, h.ζsP, h.logσ_ζP, approx, rnormPM.P, ϕqPc, cor_ends.P, h.sample_buffers, ) (;dϕqPc = gradh.dϕqPc, dϕqIc = gradh.dϕqIc, dϕg = gradh.dϕg), gradh end +function set_∂elbo_∂logσ_ζP!(∂elbo_∂logσ_ζP::AbstractVector{T}, ::NegElbo) where T + ∂elbo_∂logσ_ζP .= -one(T) +end +function set_∂elbo_∂logσ_ζP!(∂elbo_∂logσ_ζP::AbstractVector{T}, ::NegLogDensity) where T + ∂elbo_∂logσ_ζP .= zero(T) +end + + +function pullack_sample_ζsP!(dϕqPc, ::NegElbo, dζsP, dlogσ_ζP, gradh, args...) + gradh.pullback_cl_sample_ζsP!(dϕqPc, dζsP, dlogσ_ζP, args...) +end +function pullack_sample_ζsP!(dϕqPc, ::NegLogDensity, dζsP, dlogσ_ζP::AbstractVector{TF}, args...) where TF + #error("debug pullack_sample_ζsP!(dϕqPc, ::NegLogDensity") + n_θM = size(dζsP, 1) + @assert all(dlogσ_ζP .== zero(TF)) + dϕqPc[1:n_θM] .= dζsP + dϕqPc[(n_θM+1):end] .= zero(TF) + dϕqPc +end + + """ Callable to make deliver arguments that do not differ by individual to Foldl.mapreduce. """ -struct ForwardDiffGradNelboiZCl{KEYS, NMC, TA, Tϕq, TD, THWC, TC, TM, TAF} +struct ForwardDiffGradNelboiZCl{KEYS, NMC, TA, Tϕq, TD, THWC, TC, TM, TAF, LM} approx::TA ϕqIc::Tϕq # TODO remove from closure, provided with preallocated VcatVecViews dϕmvecs::TD @@ -128,6 +157,7 @@ struct ForwardDiffGradNelboiZCl{KEYS, NMC, TA, Tϕq, TD, THWC, TC, TM, TAF} # need to store KEYS in type parameter otherwise allocation in NamedTuple(keys) keys_nljoint_args_indi::Val{KEYS} nljoint_args_fix::TAF + loss_mode::LM n_MC_val::Val{NMC} end function (f::ForwardDiffGradNelboiZCl{KEYS})(tup) where KEYS @@ -152,7 +182,7 @@ function forwarddiff_grad_nelboi_z!(hi, rnormM, hwi, hwi_dc = hwi_tup #grad_conf = hwi.grad_conf nelboi_vv = make_nelboiz_lazy_cls_vv(hi, cl.approx, rnormM, nljoint_args_indi, - cl.cor_endsM, cl.transM, cl.nljoint_args_fix, ax_ϕqI) + cl.cor_endsM, cl.transM, cl.nljoint_args_fix, cl.loss_mode, ax_ϕqI) # grads are vector-shaped, but copyto! dϕmvecs and SAVector use those anyhow grad_θsP, grad_ϕm, grad_ϕqI = if isnothing(omit_gradient) # θsP_ϕm_ϕqI is preallocated per worker (hwi.θsP_ϕm_ϕqI). @@ -184,6 +214,7 @@ end function make_nelboiz_lazy_cls_vv(hi, approx::AbstractHVIApproximation, rnormM, nljoint_args_indi, cor_endsM, transM, nljoint_args_fix, + loss_mode::AbstractElboLossMode, ax_ϕqI::CA.Axis ) # VcatVecViews(θsP, ϕm, ϕqI) @@ -194,7 +225,7 @@ function make_nelboiz_lazy_cls_vv(hi, approx::AbstractHVIApproximation, ϕm_dual = subarray(vv,Val(2)) ϕqIc_dual = CA.ComponentVector(subarray(vv,Val(3)), ax_ϕqI) compute_nelboi_z!( - hi, approx, rnormM, nljoint_args_indi, + hi, loss_mode, approx, rnormM, nljoint_args_indi, ϕm_dual, ϕqIc_dual, θsP_dual, cor_endsM, transM, nljoint_args_fix, )[1] @@ -369,9 +400,16 @@ function prepare_gradelbo_helpers( n_site, n_cov, cor_ends, transP::Stacked, transM::Stacked, nljoint_args_fix, - n_MC_val::Val{NMC}, + loss_mode::AbstractElboLossMode, + n_MC_val::Val{NMC_ELBO}, max_diffchunk::ForwardDiff.Chunk{chunk} = ForwardDiff.Chunk(8), - ) where {TG, TF, ND, NMC, chunk} + ) where {TG, TF, ND, NMC_ELBO, chunk} + if loss_mode isa NegLogDensity + NMC = 1 + n_MC_val = Val(1) + else + NMC = NMC_ELBO + end local hi1, ax_θP, ax_θM, n_covP, n_ϕmvec, n_θP, n_MC, n_M hi1 = h.helpers_sites[1] ax_θP = get_axis_θP(nljoint_args_fix.f!) @@ -401,7 +439,7 @@ function prepare_gradelbo_helpers( # local nelboiz_vv, θsP_ϕm_ϕqI_vv1, diffchunks_elbo_vv nelboiz_vv = make_nelboiz_lazy_cls_vv(hi1, approx, rnormMi, nljoint_args_indi, - cor_ends.M, transM, nljoint_args_fix, CA.getaxes(ϕqIc)[1], ) + cor_ends.M, transM, nljoint_args_fix, loss_mode, CA.getaxes(ϕqIc)[1], ) θsP_ϕm_ϕqI_vv1 = θsP_ϕm_ϕqI_vv(θsP,ϕm,ϕqIc) # n_M >= n_θM but not static diffchunks_elbo_vv = ForwardDiff.Chunk(min(chunk, diff --git a/test/test_elbo_site.jl b/test/test_elbo_site.jl index 1f2d489..53a137a 100644 --- a/test/test_elbo_site.jl +++ b/test/test_elbo_site.jl @@ -35,7 +35,7 @@ n_covP2 = 2 n_cov = 3 n_θP = 3 n_θM = 3 -n_M = n_θM + 1 # additional uncertainty scaling factor +n_M = n_θM + 1 # additional uncertainty scaling factor for ScalineApprox import Lux import Zygote @@ -171,6 +171,7 @@ isUsingSimpleChains = false # preallocate helpers and shadows rnormPM = CP.prepare_rnorm(ϕqP; n_θP, n_θM, n_MC, n_site) + rnormPM_point = CP.prepare_rnorm(ϕqP; n_θP, n_θM, n_MC=1, n_site) # size(rnormPM.P) approx = CP.DiagonalHVIApproximation() CP.randnPM!(rng, rnormPM) @@ -180,6 +181,12 @@ isUsingSimpleChains = false CP.check_elbo_helpers(h2, xM, pbm_covar_indices2; n_ϕg = length(ϕg2)) randn!(h0.θsP) # for testing should initialized to finite values randn!(h2.θsP) # for testing should initialized to finite values + h0_point = CP.prepare_elbo_helpers(approx, ϕg, ϕqP; n_θP, n_θM, n_site, n_MC=1, n_cov, n_covP = n_covP0, n_M, cor_ends, n_obs) + h2_point = CP.prepare_elbo_helpers(approx, ϕg2, ϕqP; n_θP, n_θM, n_site, n_MC=1, n_cov, n_covP = n_covP2, n_M, cor_ends, n_obs) + CP.check_elbo_helpers(h0_point, xM, pbm_covar_indices0; n_ϕg = length(ϕg)) + CP.check_elbo_helpers(h2_point, xM, pbm_covar_indices2; n_ϕg = length(ϕg2)) + randn!(h0_point.θsP) # for testing should initialized to finite values + randn!(h2_point.θsP) # for testing should initialized to finite values # approxM = CP.MeanHVIApproximation() h0M = CP.prepare_elbo_helpers(approxM, ϕg, ϕqP; n_θP, n_θM, n_site, n_MC, n_cov, n_covP = n_covP0, n_M, cor_ends, n_obs) @@ -192,8 +199,10 @@ isUsingSimpleChains = false approxS = CP.MeanUniScalingHVIApproximation(ϕqIc1.logσ_ζM[1]) h0S = CP.prepare_elbo_helpers(approxS, ϕg, ϕqP; n_θP, n_θM, n_site, n_MC, n_cov, n_covP = n_covP0, n_M, cor_ends, n_obs) h2S = CP.prepare_elbo_helpers(approxS, ϕg2, ϕqP; n_θP, n_θM, n_site, n_MC, n_cov, n_covP = n_covP2, n_M, cor_ends, n_obs) + h2S_point = CP.prepare_elbo_helpers(approxS, ϕg2, ϕqP; n_θP, n_θM, n_site, n_MC=1, n_cov, n_covP = n_covP2, n_M, cor_ends, n_obs) randn!(h0S.θsP) # for testing should initialized to finite values randn!(h2S.θsP) # for testing should initialized to finite values + randn!(h2S_point.θsP) # for testing should initialized to finite values #------------------------------------------------------------- @testset "_setU_scaled!" begin @@ -398,6 +407,7 @@ end ϕqIc = intϕqI(ϕqI) ϕqPc = intϕqP(ϕqP) @assert keys(indiv_args1)[1] == :xM + loss_mode = CP.NegElbo() xM = indiv_args1.xM; r2end = 2:length(indiv_args1) nljoint_args_inds = NamedTuple{keys(indiv_args1)[r2end]}(values(indiv_args1)[r2end]) CP.sample_ζsP!(h0.ζsP, h0.logσ_ζP, approx, rnormPM.P, ϕqPc, cor_ends.P, h0.sample_buffers) # n_P * n_MC @@ -408,48 +418,49 @@ end rnormM1 = rnormPM.M[1] ϕms1 = h0.ϕms[:,1] θsP1 = h0.θsP - CP.compute_nelboi_z!(hi1, approx, rnormM1, nljoint_args_ind1, ϕms1, ϕqIc, θsP1, + CP.compute_nelboi_z!(hi1, CP.NegElbo(), approx, rnormM1, nljoint_args_ind1, ϕms1, ϕqIc, θsP1, cor_ends.M, transM, nljoint_args_fix) #@code_warntype CP.compute_nelboi_z!(hi1, rnormM1, nljoint_args_ind1, ϕms1, ϕqIc, θsP1) # need two wrap in tmpf to avoid allocations due to boxing - function tmpf(hi, approx, rnormM1, nljoint_args_ind1, ϕms1, ϕqIc, θsP1, cor_endsM, transM, nljoint_args_fix) - @test (@allocated CP.compute_nelboi_z!(hi, approx, rnormM1, nljoint_args_ind1, ϕms1, + function tmpf(hi, loss_mode, approx, rnormM1, nljoint_args_ind1, ϕms1, ϕqIc, θsP1, cor_endsM, transM, nljoint_args_fix) + @test (@allocated CP.compute_nelboi_z!(hi, loss_mode, approx, rnormM1, nljoint_args_ind1, ϕms1, ϕqIc, θsP1, cor_endsM, transM, nljoint_args_fix)) == 0 end - tmpf(hi1, approx, rnormM1, nljoint_args_ind1, ϕms1, ϕqIc, θsP1, cor_ends.M, transM, nljoint_args_fix) - #@profview CP.compute_nelboi_z!(hi1, approx, rnormM1, nljoint_args_ind1, ϕms1, ϕqIc, θsP1, cor_ends.M, transM, nljoint_args_fix) - #using BenchmarkTools - #@btime CP.compute_nelboi_z!($hi1, $approx, $rnormM1, $nljoint_args_ind1, $ϕms1, $ϕqIc, $θsP1) + tmpf(hi1, CP.NegElbo(), approx, rnormM1, nljoint_args_ind1, ϕms1, ϕqIc, θsP1, cor_ends.M, transM, nljoint_args_fix) + CP.compute_nelboi_z!(hi1, CP.NegLogDensity(), approx, rnormM1, nljoint_args_ind1, ϕms1, ϕqIc, θsP1, + cor_ends.M, transM, nljoint_args_fix) + tmpf(hi1, CP.NegLogDensity(), approx, rnormM1, nljoint_args_ind1, ϕms1, ϕqIc, θsP1, cor_ends.M, transM, nljoint_args_fix) # # test with views as input to compute_nelboi_z! and differnt ϕm per MC h21 = h2.helpers_sites[1] CP.g_apply!(h2.ϕms_mcs, ϕg2, xM, h2.ζsP, pbm_covar_indices2, g2, h2.xMP, false) ϕm = randn!(similar(h2.ϕms_mcs[:,:,1])) inputs = CA.ComponentVector(ϕqIc=ϕqIc, ϕm = ϕm, θsP= θsP1) - CP.compute_nelboi_z!(h21, approx, rnormM1, nljoint_args_ind1, inputs.ϕm, inputs.ϕqIc, + CP.compute_nelboi_z!(h21, CP.NegElbo(), approx, rnormM1, nljoint_args_ind1, inputs.ϕm, inputs.ϕqIc, inputs.θsP, cor_ends.M, transM, nljoint_args_fix) # with ϕms1_ = view(inputs, Val(:ϕm)) ϕqIc_ = view(inputs, Val(:ϕqIc)) θsP1_ = view(inputs, Val(:θsP)) - function loop_compute_nelboi_z(n, h21, approx, rnormM1, nljoint_args_ind1, ϕms1_, ϕqIc_, θsP1_, cor_endsM, transM, nljoint_args_fix) + function loop_compute_nelboi_z(n, h21, loss_mode, approx, rnormM1, nljoint_args_ind1, ϕms1_, ϕqIc_, θsP1_, cor_endsM, transM, nljoint_args_fix) for _ in 1:n - @noinline CP.compute_nelboi_z!(h21, approx, rnormM1, nljoint_args_ind1, ϕms1_, ϕqIc_, θsP1_, cor_endsM, transM, nljoint_args_fix) + @noinline CP.compute_nelboi_z!(h21, loss_mode, approx, rnormM1, nljoint_args_ind1, ϕms1_, ϕqIc_, θsP1_, cor_endsM, transM, nljoint_args_fix) end nothing end - function alloc_compute_nelboi_z(h21, approx, rnormM1, nljoint_args_ind1, ϕms1_, ϕqIc_, θsP1_, cor_endsM, transM, nljoint_args_fix) + function alloc_compute_nelboi_z(h21, loss_mode, approx, rnormM1, nljoint_args_ind1, ϕms1_, ϕqIc_, θsP1_, cor_endsM, transM, nljoint_args_fix) # avoid global variables -> pass them through function - loop_compute_nelboi_z(1, h21, approx, rnormM1, nljoint_args_ind1, ϕms1_, ϕqIc_, θsP1_, cor_endsM, transM, nljoint_args_fix) # thorough warmup - #@profview_allocs loop_compute_nelboi_z(10_000,h21, approx, rnormM1, nljoint_args_ind1, ϕms1_, ϕqIc_, θsP1_, cor_endsM, transM, nljoint_args_fix) - @test (@allocated loop_compute_nelboi_z(100,h21, approx, rnormM1, nljoint_args_ind1, ϕms1_, ϕqIc_, θsP1_, cor_endsM, transM, nljoint_args_fix)) == 0 + loop_compute_nelboi_z(1, h21, loss_mode, approx, rnormM1, nljoint_args_ind1, ϕms1_, ϕqIc_, θsP1_, cor_endsM, transM, nljoint_args_fix) # thorough warmup + #@profview_allocs loop_compute_nelboi_z(10_000,h21, loss_mode, approx, rnormM1, nljoint_args_ind1, ϕms1_, ϕqIc_, θsP1_, cor_endsM, transM, nljoint_args_fix) + @test (@allocated loop_compute_nelboi_z(100,h21, loss_mode, approx, rnormM1, nljoint_args_ind1, ϕms1_, ϕqIc_, θsP1_, cor_endsM, transM, nljoint_args_fix)) == 0 end - alloc_compute_nelboi_z(h21, approx, rnormM1, nljoint_args_ind1, ϕms1_, ϕqIc_, θsP1_, cor_ends.M, transM, nljoint_args_fix) + alloc_compute_nelboi_z(h21, CP.NegElbo(), approx, rnormM1, nljoint_args_ind1, ϕms1_, ϕqIc_, θsP1_, cor_ends.M, transM, nljoint_args_fix) + alloc_compute_nelboi_z(h21, CP.NegLogDensity(), approx, rnormM1, nljoint_args_ind1, ϕms1_, ϕqIc_, θsP1_, cor_ends.M, transM, nljoint_args_fix) # gradh2 = @inferred CP.prepare_gradelbo_helpers(inputs.ϕqIc, inputs.ϕm, inputs.θsP, ϕg, ϕqPc, approx; pbm_covar_indices=pbm_covar_indices2, n_workers=1, h=h2, rnormMi = rnormPM.M[1], nljoint_args_indi=nljoint_args_ind1, - n_site, n_cov, cor_ends, transP, transM, nljoint_args_fix, n_MC_val) + n_site, n_cov, cor_ends, transP, transM, nljoint_args_fix, loss_mode, n_MC_val) hw_channel = gradh2.hw_channel # alternative: invoke forwarddiff_grad_nelboi_z! with ϕqIc_ and θsP_ as views, # passing dϕmvecs as the plain array. The views are created once *outside* the @@ -467,7 +478,7 @@ end dθsP = gradh2.dθsP cl = HVI.ForwardDiffGradNelboiZCl( approx, ϕqIc, gradh2.dϕmvecs, gradh2.hw_channel, cor_ends.M, transM, - Val(keys(nljoint_args_ind1)), nljoint_args_fix, n_MC_val) + Val(keys(nljoint_args_ind1)), nljoint_args_fix, loss_mode, n_MC_val) function alloc_forwarddiff_grad_nelboi_z(hi, rnormM, nljoint_args_indi, ϕm_, i, cl) loop_forwarddiff_grad_nelboi_z(1, hi, rnormM, nljoint_args_indi, ϕm_, i, cl, true) @test (@allocated loop_forwarddiff_grad_nelboi_z(100, hi, rnormM, nljoint_args_indi, ϕm_, i, cl, true)) == 0 @@ -485,11 +496,11 @@ end approx2 = MeanHVIApproximation() hiM1 = h0M.helpers_sites[1] - @test isfinite(CP.compute_nelboi_z!(hiM1, approx2, rnormM1, nljoint_args_ind1, + @test isfinite(CP.compute_nelboi_z!(hiM1, CP.NegElbo(), approx2, rnormM1, nljoint_args_ind1, ϕms1, ϕqIc, θsP1, cor_ends.M, transM, nljoint_args_fix)) function alloc_compute_nelboi_zM!(hiM1, approx2, rnormM1, nljoint_args_ind1, ϕms1, ϕqIc, θsP1, cor_endsM, transM, nljoint_args_fix) - @test (@allocated CP.compute_nelboi_z!(hiM1, approx2, rnormM1, nljoint_args_ind1, + @test (@allocated CP.compute_nelboi_z!(hiM1, CP.NegElbo(), approx2, rnormM1, nljoint_args_ind1, ϕms1, ϕqIc, θsP1, cor_endsM, transM, nljoint_args_fix)) == 0 end alloc_compute_nelboi_zM!(hiM1, approx2, rnormM1, nljoint_args_ind1, @@ -658,7 +669,10 @@ end # to be differentiated, return only first entry of result function _ftmp2(h0, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix) - CP.neg_elbo_sites!(h0, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix)[1] + CP.neg_elbo_sites!(h0, CP.NegElbo(), ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix)[1] +end +function _ftmp_point(h0, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix) + CP.neg_elbo_sites!(h0, CP.NegLogDensity(), ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix)[1] end function grad_neg_elbo_sites_enzyme() # differentiate entire neg_elbo_sites by enzyme @@ -674,7 +688,7 @@ function grad_neg_elbo_sites_enzyme() # differentiate entire neg_elbo_sites by e ϕ = (;ϕg=ϕgv, ϕqP, ϕqI) dϕ = map(zero, ϕ) sample_args = (; - approx, g, is_testmode, pbm_covar_indices = nothing, intϕqP, intϕqI, cor_ends, + approx, g, is_testmode = false, pbm_covar_indices = nothing, intϕqP, intϕqI, cor_ends, transP, transM) # rng1 = StableRNG(1234) @@ -696,9 +710,9 @@ function grad_neg_elbo_sites_enzyme() # differentiate entire neg_elbo_sites by e Enzyme.Const(nljoint_args_fix), ) () -> begin - dϕg0_enz ≈ dϕ.ϕg - dϕqP0_enz ≈ dϕ.ϕqP - dϕqI0_enz ≈ dϕ.ϕqI + @test dϕg0_enz ≈ dϕ.ϕg + @test dϕqP0_enz ≈ dϕ.ϕqP + @test dϕqI0_enz ≈ dϕ.ϕqI hcat(dϕg0_enz, dϕ.ϕg) end dϕg0_enz = copy(dϕ.ϕg) @@ -722,7 +736,7 @@ function grad_neg_elbo_sites_enzyme() # differentiate entire neg_elbo_sites by e ϕ2 = (;ϕg=ϕg2v, ϕqP, ϕqI) dϕ2 = map(zero, ϕ2) sample_args2 = (; - approx, g=g2, is_testmode, pbm_covar_indices = pbm_covar_indices2, intϕqP, intϕqI, cor_ends, + approx, g=g2, is_testmode = false, pbm_covar_indices = pbm_covar_indices2, intϕqP, intϕqI, cor_ends, transP, transM) # rng1 = StableRNG(1234) @@ -744,10 +758,10 @@ function grad_neg_elbo_sites_enzyme() # differentiate entire neg_elbo_sites by e Enzyme.Const(nljoint_args_fix), ) () -> begin - primal2_enz ≈ primal_enz - dϕg2_enz ≈ dϕ2.ϕg + @test primal2_enz ≈ primal_enz + @test dϕg2_enz ≈ dϕ2.ϕg @test dϕqP2_enz ≈ dϕ2.ϕqP atol=0.001 - dϕqI2_enz ≈ dϕ2.ϕqI + @test dϕqI2_enz ≈ dϕ2.ϕqI hcat(dϕg2_enz, dϕ2.ϕg) hcat(dϕqP2_enz, dϕ2.ϕqP) end @@ -775,14 +789,19 @@ function grad_neg_elbo_sites_enzyme_scaled() intϕqP, intϕqI=intϕqIS, cor_ends, transP, transM) h2Sp = CP.prepare_elbo_helpers(sample_args2S.approx, ϕg2, ϕqP; - n_θP, n_θM, n_site, n_MC, n_cov, n_covP = n_covP2, n_M, cor_ends, use_diff_cache = Val(false)) + n_θP, n_θM, n_site, n_MC, n_cov, n_covP = n_covP2, n_M, cor_ends, n_obs, + use_diff_cache = Val(false)) dh2Sp = Enzyme.make_zero(h2Sp) + h2Sp_point = CP.prepare_elbo_helpers(sample_args2S.approx, ϕg2, ϕqP; + n_θP, n_θM, n_site, n_MC=1, n_cov, n_covP = n_covP2, n_M, cor_ends, n_obs, + use_diff_cache = Val(false)) + dh2Sp_point = Enzyme.make_zero(h2Sp_point) # rng1 = StableRNG(1234) CP.randnPM!(rng1, rnormPM) randn!(rng1, ϕ2S.ϕg) randn!(rng1, indiv_args.xM) - primal_enz2S = tmp = _ftmp2(h2Sp, ϕ2S, rnormPM, sample_args2S, indiv_args) + primal_enz2S = tmp = _ftmp2(h2Sp, ϕ2S, rnormPM, sample_args2S, indiv_args, nljoint_args_fix) Enzyme.make_zero!(dϕ2S) Enzyme.make_zero!(dh2Sp) Enzyme.autodiff( @@ -794,7 +813,14 @@ function grad_neg_elbo_sites_enzyme_scaled() Enzyme.DuplicatedNoNeed(rnormPM, Enzyme.make_zero(rnormPM)), Enzyme.Const(sample_args2S), Enzyme.Const(indiv_args), + Enzyme.Const(nljoint_args_fix), ) + () -> begin + @test dϕg2S_enz ≈ dϕ2S.ϕg atol = 0.03 + @test dϕqP2S_enz ≈ dϕ2S.ϕqP atol = 0.03 + @test dϕqI2S_enz ≈ dϕ2S.ϕqI atol = 0.03 + hcat(dϕg2S_enz, dϕ2S.ϕg) + end dϕg2S_enz = copy(dϕ2S.ϕg) dϕqP2S_enz = copy(dϕ2S.ϕqP) dϕqI2S_enz = copy(dϕ2S.ϕqI) @@ -806,6 +832,36 @@ function grad_neg_elbo_sites_enzyme_scaled() primal_enz2S, dϕg2S_enz, dϕqP2S_enz, dϕqI2S_enz = JLD2.load(fname, "primal_enz2S", "dϕg2S_enz", "dϕqP2S_enz", "dϕqI2S_enz"); end + # + rng1 = StableRNG(1234) + CP.randnPM!(rng1, rnormPM_point) + randn!(rng1, ϕ2S.ϕg) + randn!(rng1, indiv_args.xM) + primal_enz2S_point = tmp = _ftmp_point(h2Sp_point, ϕ2S, rnormPM_point, sample_args2S, indiv_args, nljoint_args_fix) + dϕ2S_point = CP.map_leaves_nt(copy, dϕ2S) + Enzyme.make_zero!(dϕ2S_point) + Enzyme.make_zero!(dh2Sp_point) + Enzyme.autodiff( + Enzyme.set_runtime_activity(Enzyme.Reverse) , + _ftmp_point, + Enzyme.Active, + Enzyme.Duplicated(h2Sp_point, dh2Sp_point), + Enzyme.Duplicated(ϕ2S, dϕ2S_point), + Enzyme.DuplicatedNoNeed(rnormPM_point, Enzyme.make_zero(rnormPM_point)), + Enzyme.Const(sample_args2S), + Enzyme.Const(indiv_args), + Enzyme.Const(nljoint_args_fix), + ) + dϕg2S_point_enz = copy(dϕ2S_point.ϕg) + dϕqP2S_point_enz = copy(dϕ2S_point.ϕqP) + dϕqI2S_point_enz = copy(dϕ2S_point.ϕqI) + () -> begin + fname = "intermediate/test_enzymeS_dphi2_point.jld2" + mkpath("intermediate") + JLD2.jldsave(fname, false, IOStream; primal_enz2S_point, dϕg2S_point_enz, dϕqP2S_point_enz, dϕqI2S_point_enz) + primal_enz2S_point, dϕg2S_point_enz, dϕqP2S_point_enz, dϕqI2S_point_enz = JLD2.load(fname, + "primal_enz2S_point", "dϕg2S_point_enz", "dϕqP2S_point_enz", "dϕqI2S_point_enz"); + end end function tmp_reference_revision() @@ -859,12 +915,11 @@ end sample_args = (; approx, g, is_testmode, pbm_covar_indices = nothing, intϕqP, intϕqI, cor_ends, transP, transM) - primal = CP.neg_elbo_sites!(h0, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix) + primal = CP.neg_elbo_sites!(h0, CP.NegElbo(), ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix) res0, gradh0 = CP.grad_neg_elbo_sites( - #@descend_code_warntype CP.neg_elbo_sites!( - h0, (;), ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix; n_MC_val) + h0, (;), CP.NegElbo(), ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix; n_MC_val) res0_, gradh0_ = CP.grad_neg_elbo_sites( # test deterministic result and distributed - h0, (;), ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix; n_MC_val, + h0, (;), CP.NegElbo(), ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix; n_MC_val, executor = distributedEx ) @test all(map(≈, res0_, res0)) @@ -883,6 +938,16 @@ end #dϕqP0_enz - CA.getdata(res0.dϕqP) end # + # pointloss + res0_point, gradh0_point = CP.grad_neg_elbo_sites( + h0_point, (;), CP.NegLogDensity(), ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix; n_MC_val) + @test all(res0_point.dϕqPc.μζP .!= zero(eltype(res0_point.dϕqPc))) + @test all(res0_point.dϕqPc.logσ_ζP .== zero(eltype(res0_point.dϕqPc))) + @test all(res0_point.dϕqPc.ρsP .== zero(eltype(res0_point.dϕqPc))) + @test all(res0_point.dϕqIc.logσ_ζM .== zero(eltype(res0_point.dϕqIc))) + @test all(res0_point.dϕqIc.ρsM .== zero(eltype(res0_point.dϕqIc))) + # + # #---------------- matrix mode with population covariates ϕ2 = (;ϕ..., ϕg=ϕg2v) sample_args2 = (;sample_args..., g=g2, pbm_covar_indices = pbm_covar_indices2) @@ -890,9 +955,9 @@ end CP.randnPM!(rng1, rnormPM) randn!(rng1, ϕ2.ϕg) randn!(rng1, indiv_args.xM) - primal2 = CP.neg_elbo_sites!(h2, ϕ2, rnormPM, sample_args2, indiv_args, nljoint_args_fix)[1] - res0, gradh2 = CP.grad_neg_elbo_sites(h2, (;), ϕ2, rnormPM, sample_args2, indiv_args, nljoint_args_fix; n_MC_val) - res0_, gradh2_ = CP.grad_neg_elbo_sites(h2, (;), ϕ2, rnormPM, sample_args2, indiv_args, nljoint_args_fix; n_MC_val, + primal2 = CP.neg_elbo_sites!(h2, CP.NegElbo(), ϕ2, rnormPM, sample_args2, indiv_args, nljoint_args_fix)[1] + res0, gradh2 = CP.grad_neg_elbo_sites(h2, (;), CP.NegElbo(), ϕ2, rnormPM, sample_args2, indiv_args, nljoint_args_fix; n_MC_val) + res0_, gradh2_ = CP.grad_neg_elbo_sites(h2, (;), CP.NegElbo(), ϕ2, rnormPM, sample_args2, indiv_args, nljoint_args_fix; n_MC_val, executor = distributedEx,) @test all(map(≈, res0_, res0)) # if we saved Enzyme results earlier, compare to them @@ -906,6 +971,14 @@ end @test dϕqP2_enz ≈ res0.dϕqPc atol = 0.001 #hcat(dϕqP2_enz, res0.dϕqP) end + res0_point, gradh2_point = CP.grad_neg_elbo_sites( + h2_point, (;), CP.NegLogDensity(), ϕ2, rnormPM, sample_args2, indiv_args, nljoint_args_fix; n_MC_val) + @test all(res0_point.dϕqPc.μζP .!= zero(eltype(res0_point.dϕqPc))) + @test all(res0_point.dϕqPc.logσ_ζP .== zero(eltype(res0_point.dϕqPc))) + @test all(res0_point.dϕqPc.ρsP .== zero(eltype(res0_point.dϕqPc))) + @test all(res0_point.dϕqIc.logσ_ζM .== zero(eltype(res0_point.dϕqIc))) + @test all(res0_point.dϕqIc.ρsM .== zero(eltype(res0_point.dϕqIc))) + #---------------- approxM with non-empty h.sample_buffers and scaling # and using eachcol in PBM function ϕ2S = (;ϕ2..., ϕqI=CA.getdata(ϕqIcS)) @@ -918,9 +991,9 @@ end CP.randnPM!(rng1, rnormPM) randn!(rng1, ϕ2S.ϕg) randn!(rng1, indiv_args.xM) - primal2S = CP.neg_elbo_sites!(h2S, ϕ2S, rnormPM, sample_args2S, indiv_args, nljoint_args_fix2S) + primal2S = CP.neg_elbo_sites!(h2S, CP.NegElbo(), ϕ2S, rnormPM, sample_args2S, indiv_args, nljoint_args_fix2S) res0, gradh2S = CP.grad_neg_elbo_sites( - h2S, (;), ϕ2S, rnormPM, sample_args2S, indiv_args, nljoint_args_fix2S; n_MC_val, + h2S, (;), CP.NegElbo(), ϕ2S, rnormPM, sample_args2S, indiv_args, nljoint_args_fix2S; n_MC_val, #executor = distributedEx, ) if isfile("intermediate/test_enzymeS_dphi2.jld2") @@ -934,26 +1007,47 @@ end #hcat(dϕqP2_enz, res0.dϕqP) #hcat(dϕqI2S_enz, res0.dϕqI) end + # + rng1 = StableRNG(1234) + CP.randnPM!(rng1, rnormPM_point) + randn!(rng1, ϕ2S.ϕg) + randn!(rng1, indiv_args.xM) + primal2S_point = CP.neg_elbo_sites!(h2S_point, CP.NegLogDensity(), ϕ2S, rnormPM_point, sample_args2S, indiv_args, nljoint_args_fix2S) + res0_point, gradh2S_point = CP.grad_neg_elbo_sites( + h2S_point, (;), CP.NegLogDensity(), ϕ2S, rnormPM_point, sample_args2S, indiv_args, nljoint_args_fix2S; n_MC_val, + #executor = distributedEx, + ) + if isfile("intermediate/test_enzymeS_dphi2_point.jld2") + primal_enz2S_point, dϕg2S_point_enz, dϕqP2S_point_enz, dϕqI2S_point_enz = JLD2.load( + "intermediate/test_enzymeS_dphi2_point.jld2", + "primal_enz2S_point", "dϕg2S_point_enz", "dϕqP2S_point_enz", "dϕqI2S_point_enz"); + @test primal_enz2S_point ≈ primal2S_point[1] + @test dϕg2S_point_enz ≈ res0_point.dϕg atol=0.03 + @test dϕqI2S_point_enz ≈ res0_point.dϕqIc atol=0.01 + @test dϕqP2S_point_enz ≈ res0_point.dϕqPc atol=0.01 + #hcat(dϕqP2_enz, res0.dϕqP) + #hcat(dϕqI2S_enz, res0.dϕqI) + end function loop_grad_neg_elbo_sites( - n, h2S, gradh2S, ϕ2S, rnormPM, sample_args2S, indiv_args, nljoint_args_fix) + n, h2S, gradh2S, loss_mode, ϕ2S, rnormPM, sample_args2S, indiv_args, nljoint_args_fix) for _ in 1:n @noinline CP.grad_neg_elbo_sites( - h2S, gradh2S, ϕ2S, rnormPM, sample_args2S, indiv_args, nljoint_args_fix; n_MC_val) + h2S, gradh2S, loss_mode, ϕ2S, rnormPM, sample_args2S, indiv_args, nljoint_args_fix; n_MC_val) end nothing end - function alloc_grad_neg_elbo_sites(h2S, gradh2S, ϕ2S, rnormPM, sample_args2S, indiv_args, nljoint_args_fix) + function alloc_grad_neg_elbo_sites(h2S, gradh2S, loss_mode, ϕ2S, rnormPM, sample_args2S, indiv_args, nljoint_args_fix) # avoid global variables -> pass them through function - loop_grad_neg_elbo_sites(1, h2S, gradh2S, ϕ2S, rnormPM, sample_args2S, indiv_args, nljoint_args_fix) - #@profview_allocs loop_grad_neg_elbo_sites(1_000,h2S, gradh2S, ϕ2S, rnormPM, sample_args2S, indiv_args, nljoint_args_fix) - @test (@allocated loop_grad_neg_elbo_sites(100,h2S, gradh2S, ϕ2S, rnormPM, sample_args2S, indiv_args, nljoint_args_fix)) <= 3_264_000 #3_260_800 + loop_grad_neg_elbo_sites(1, h2S, gradh2S, loss_mode, ϕ2S, rnormPM, sample_args2S, indiv_args, nljoint_args_fix) + #@profview_allocs loop_grad_neg_elbo_sites(1_000,h2S, gradh2S, loss_mode, ϕ2S, rnormPM, sample_args2S, indiv_args, nljoint_args_fix) + @test (@allocated loop_grad_neg_elbo_sites(100,h2S, gradh2S, loss_mode, ϕ2S, rnormPM, sample_args2S, indiv_args, nljoint_args_fix)) <= 3_264_000 #3_260_800 end - alloc_grad_neg_elbo_sites(h2S, gradh2S, ϕ2S, rnormPM, sample_args2S, indiv_args, nljoint_args_fix2S) + alloc_grad_neg_elbo_sites(h2S, gradh2S, CP.NegElbo(), ϕ2S, rnormPM, sample_args2S, indiv_args, nljoint_args_fix2S) #_i_sites_train = 1:n_site #@usingany BenchmarkTools #@benchmark CP.grad_neg_elbo_sites($h2, $gradh2, $rnormPM, $ϕg2v, $ϕqP, $ϕqI, $g2, $pbm_covar_indices2; i_sites_train = _i_sites_train, intϕqP=$intϕqP, intϕqI=$intϕqI, xM=$xM, is_testmode = false) - + alloc_grad_neg_elbo_sites(h2S_point, gradh2S_point, CP.NegLogDensity(), ϕ2S, rnormPM, sample_args2S, indiv_args, nljoint_args_fix2S) end diff --git a/test/test_elbo_site_doubleMM.jl b/test/test_elbo_site_doubleMM.jl index 0860871..ff40d8b 100644 --- a/test/test_elbo_site_doubleMM.jl +++ b/test/test_elbo_site_doubleMM.jl @@ -84,7 +84,9 @@ pbm_covar_indices = SA.SVector{length(pbm_covars)}([findfirst(==(k), keys(par_te # approxS = get_hybridproblem_HVIApproximation(prob; scenario) n_M = get_numberof_MLinputs(approx, par_templates.θM) h0 = CP.prepare_elbo_helpers(approx, ϕg, ϕqP; n_θP, n_θM, n_site, n_MC, n_cov, n_covP, n_M, cor_ends, n_obs); + h0_point = CP.prepare_elbo_helpers(approx, ϕg, ϕqP; n_θP, n_θM, n_site, n_MC=1, n_cov, n_covP, n_M, cor_ends, n_obs); randn!(h0.θsP) # for testing should initialized to finite values + randn!(h0_point.θsP) # for testing should initialized to finite values rnormPM = CP.prepare_rnorm(ϕqP; n_θP, n_θM, n_MC, n_site) #------------------------------------------------------------- @@ -102,10 +104,10 @@ pbm_covar_indices = SA.SVector{length(pbm_covars)}([findfirst(==(k), keys(par_te sample_args = (; approx, g, is_testmode = true, pbm_covar_indices, intϕqP, intϕqI, cor_ends, transP, transM) - primal = CP.neg_elbo_sites!(h0, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix) + primal = CP.neg_elbo_sites!(h0, CP.NegElbo(), ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix) function loop_neg_elbo_sites!(n, h0, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix) for _ in 1:n - @noinline CP.neg_elbo_sites!(h0, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix) + @noinline CP.neg_elbo_sites!(h0, CP.NegElbo(), ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix) end nothing end @@ -120,32 +122,36 @@ pbm_covar_indices = SA.SVector{length(pbm_covars)}([findfirst(==(k), keys(par_te local sample_args_rev, res0, gradh0 sample_args_rev = (; sample_args..., is_testmode=false) res0, gradh0 = @inferred CP.grad_neg_elbo_sites( - #@descend_code_warntype CP.neg_elbo_sites!( - h0, (;), ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix; n_MC_val, + h0, (;), CP.NegElbo(), ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix; n_MC_val, + #n_workers = 3, + ) + res0_point, gradh0_point = @inferred CP.grad_neg_elbo_sites( + h0_point, (;), CP.NegLogDensity(), ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix; n_MC_val, #n_workers = 3, ) function loop_grad_neg_elbo_sites( - n, h0, gradh0, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix; n_MC_val) + n, h0, gradh0, loss_mode, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix; n_MC_val) for _ in 1:n @noinline CP.grad_neg_elbo_sites( - h0, gradh0, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix; n_MC_val) + h0, gradh0, loss_mode, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix; n_MC_val) end nothing end - function alloc_grad_neg_elbo_sites(h0, gradh0, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix; n_MC_val) + function alloc_grad_neg_elbo_sites(h0, gradh0, loss_mode, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix; n_MC_val) # avoid global variables -> pass them through function - loop_grad_neg_elbo_sites(1, h0, gradh0, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix; n_MC_val) + loop_grad_neg_elbo_sites(1, h0, gradh0, loss_mode, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix; n_MC_val) GC.gc() - #@profview_allocs loop_grad_neg_elbo_sites(10_000,h0, gradh0, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix; n_MC_val) - #@profview loop_grad_neg_elbo_sites(10_000,h0, gradh0, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix; n_MC_val) - #@test (@allocated loop_grad_neg_elbo_sites(100,h0, gradh0, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix; n_MC_val)) <= 6_748_800 + #@profview_allocs loop_grad_neg_elbo_sites(10_000,h0, gradh0, loss_mode, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix; n_MC_val) + #@profview loop_grad_neg_elbo_sites(10_000,h0, gradh0, loss_mode, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix; n_MC_val) + #@test (@allocated loop_grad_neg_elbo_sites(100,h0, gradh0, loss_mode, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix; n_MC_val)) <= 6_748_800 # TODO inspect why more allocations when calling from runtests or command line local alloc - alloc = @allocated loop_grad_neg_elbo_sites(100,h0, gradh0, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix; n_MC_val) + alloc = @allocated loop_grad_neg_elbo_sites(100,h0, gradh0, loss_mode, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix; n_MC_val) @test alloc <= 6_753_536 #6_748_800 # alloc end - alloc_grad_neg_elbo_sites(h0, gradh0, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix; n_MC_val) + alloc_grad_neg_elbo_sites(h0, gradh0, CP.NegElbo(), ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix; n_MC_val) + alloc_grad_neg_elbo_sites(h0_point, gradh0_point, CP.NegLogDensity(), ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix; n_MC_val) end # @testset "HybridPointSolverSite" begin