From eba9e5428621e8bb614a5757f32aee3b5a2553a2 Mon Sep 17 00:00:00 2001 From: Thomas Wutzler Date: Tue, 20 Oct 2026 16:44:41 +0000 Subject: [PATCH 1/4] work on PointSolver Site --- src/HybridSolver_site.jl | 374 ++++++++++++++++++++++++++++++++ src/elbo_site.jl | 39 +++- src/pointloss_site.jl | 207 ++++++++++++++++++ test/test_elbo_site_doubleMM.jl | 3 + 4 files changed, 620 insertions(+), 3 deletions(-) create mode 100644 src/HybridSolver_site.jl create mode 100644 src/pointloss_site.jl diff --git a/src/HybridSolver_site.jl b/src/HybridSolver_site.jl new file mode 100644 index 0000000..03c04ff --- /dev/null +++ b/src/HybridSolver_site.jl @@ -0,0 +1,374 @@ +""" +A Solver that solves for the mode of hybrid variational inference problem +([`AbstractHybridProblem`](@ref)). + + +It does not perform the full posterior inference, +but only finds the optimal parameters of the variational distribution. +This is useful for initializing the full posterior inference. + +- `alg`: The optimization algorithm to use for solving the problem, e.g. `Adam()` +""" +struct HybridPointSolverSite{A} <: AbstractHybridSolver + alg::A +end + +HybridPointSolverSite(; alg) = HybridPointSolverSite(alg) + +function CommonSolve.solve(prob::AbstractHybridProblem, solver::HybridPointSolverSite; + scenario=Val(()), rng=Random.default_rng(), + gdevs = nothing, # get_gdev_MP(scenario) + is_inferred::Val{is_infer} = Val(false), + ad_backend_loss = AutoEnzyme(), + epochs, + epochs_callback = max(1, floor(Int, epochs/10)), + callback = (state, loss_val) -> false, + is_omitting_NaNbatches = false, + is_omit_priors::Val{omit_priors} = Val(false), + clusters::AbstractVector{<:Integer} = + 1:first(get_hybridproblem_n_site_and_batch(prob; scenario)), + cluster_rep = 1, + kwargs... +) where {is_infer, omit_priors} + gdevs = isnothing(gdevs) ? get_gdev_MP(scenario) : gdevs + pt = get_hybridproblem_par_templates(prob; scenario) + g, ϕg0 = get_hybridproblem_MLapplicator(prob; scenario) + # TODO: separate parameters from problem description - right now optimized + ϕq0 = get_hybridproblem_ϕq(prob; scenario) + (; transP, transM) = get_hybridproblem_transforms(prob; scenario) + n_site, n_batch = get_hybridproblem_n_site_and_batch(prob; scenario) + ranef_spec = get_hybridproblem_ranef(prob; scenario) + ranef = get_ranef_computer( + ranef_spec, keys(pt.θM), n_site, one(eltype(ϕq0))) + ϕq_ranef = setup_ϕq_ranef(ranef) + if :ranef ∉ keys(ϕq0) + ϕq0 = CA.ComponentVector(ϕq0, ranef = ϕq_ranef) + else + @assert size(ϕq0[Val(:ranef)]) == size(ϕq_ranef) + end + ϕP0 = ϕq0[Val(:μP)] + #intϕ = ComponentArrayInterpreter(CA.ComponentVector(ϕg=1:length(ϕg0), ϕq=ϕq0)) + #ϕ0_cpu = vcat(ϕg0, pt.θP .* FT(0.9)) # slightly disturb θP_true + ϕ0_cpu = CA.ComponentVector(ϕg=ϕg0, ϕq=ϕq0) + intϕ = ComponentArrayInterpreter(ϕ0_cpu) + n_sites_cluster = [count(==(element),clusters) for element in 1:maximum(clusters)] + frac_cluster_all = (1 / cluster_rep) ./ n_sites_cluster[clusters] + train_loader = get_hybridproblem_train_dataloader(prob; scenario) + test_data = get_hybridproblem_test_data(prob; scenario) + # i_test = rand(1:n_site, Integer(floor(n_site/10))) + # test_data = map(train_loader.data) do data_comp + # ndims(data_comp) == 2 ? data_comp[:, i_test] : data_comp[i_test] + # end + gdev = gdevs.gdev_M + if gdev isa MLDataDevices.AbstractGPUDevice + ϕ0_dev = gdev(ϕ0_cpu) + g_dev = gdev(g) + train_loader_dev = gdev_hybridproblem_dataloader(train_loader; gdevs) + test_data_dev = gdev_hybridproblem_data( + test_data[keys(test_data)[1:5]]; gdevs) + else + ϕ0_dev = ϕ0_cpu + g_dev = g + train_loader_dev = train_loader + test_data_dev = test_data[keys(test_data)[1:5]] + end + f = get_hybridproblem_PBmodel(prob; scenario) + ftest = create_nsite_applicator(f, size(test_data[1],2)) + if gdevs.gdev_P isa MLDataDevices.AbstractGPUDevice + f_dev = gdevs.gdev_P(f) + ftest_dev = gdevs.gdev_P(ftest) + else + f_dev = f + ftest_dev = ftest + end + py = get_hybridproblem_neg_logden_obs(prob; scenario) + pbm_covars = get_hybridproblem_pbmpar_covars(prob; scenario) + n_site_test = size(test_data[1],2) + priors = get_hybridproblem_priors(prob; scenario) + priorsP = Tuple(priors[k] for k in keys(pt.θP)) + priorsM = Tuple(priors[k] for k in keys(pt.θM)) + #intP = ComponentArrayInterpreter(pt.θP) + # intθP = ComponentArrayInterpreter(pt.θP) + # intθM = ComponentArrayInterpreter(pt.θM) + penalty_computer = get_hybridproblem_penalty_computer(prob; scenario) + # loss_gf = get_loss_gf(g_dev, transM, transP, f_dev, py, intϕ; + # n_site_batch=n_batch, + # par_templates = pt, + # cdev=infer_cdev(gdevs), pbm_covars, + # priorsP, priorsM, is_omit_priors, penalty_computer, + # #intθM, intθP, + # frac_cluster_all, ranef, + # ) + # loss_gf_test = get_loss_gf(g_dev, transM, transP, ftest_dev, py, intϕ; + # n_site_batch=n_site_test, + # par_templates = pt, + # cdev=infer_cdev(gdevs), pbm_covars, + # priorsP, priorsM, is_omit_priors, penalty_computer, + # #intθM, intθP, + # frac_cluster_all, ranef, + # ) + # # call loss function once + # l1 = is_infer ? + # Test.@inferred(loss_gf(ϕ0_dev, first(train_loader_dev)...; is_testmode=true))[1] : + # # using ShareAdd; @usingany Cthulhu + # # @descend_code_warntype loss_gf(ϕ0_dev, first(train_loader_dev)...) + # loss_gf(ϕ0_dev, first(train_loader_dev)...; is_testmode=true)[1] + # l1t = loss_gf_test(ϕ0_dev, test_data_dev...; is_testmode=true)[1] + # and gradient + # xMg, xP, y_o, y_unc = first(train_loader_dev) + # gr1 = Zygote.gradient( + # p -> loss_gf(p, xMg, xP, y_o, y_unc)[1], + # ϕ0_dev) + # Zygote.gradient(ϕ0_dev -> loss_gf(ϕ0_dev, data1...)[1], ϕ0_dev) + # + # implement training loop by hand to skip minibatches with NaN gradients + ps = CA.getdata(ϕ0_dev) + opt_st_new = Optimisers.setup(solver.alg, ps) + n_skips = 0 + # prepare DI.gradient, need to access and update outside cope data_batch + # because cannot redefine fopt_loss_gf + data_batch = first(train_loader_dev) + is_testmode = false + function fopt_loss_gf(ϕ) + #@show first(data_batch[5], 2) + loss_gf(ϕ, data_batch...; is_testmode)[1] + end + ad_prep = DI.prepare_gradient(fopt_loss_gf, ad_backend_loss, zero(ps)) + grad = similar(ps) + stime = time() + for epoch in 1:epochs + is_testmode = false + #i,data_batch = first(enumerate(loader)) + for (i, data_batch_) in enumerate(train_loader_dev) + data_batch = data_batch_ # propagate outside for to scope of fopt_loss_gf + DI.gradient!(fopt_loss_gf, grad, ad_prep, ad_backend_loss, ps) + if any(isnan.(grad)) + n_skips += 1 + #println("Skipped NaN : Batch $i") + print(",$i") + else + Optimisers.update!(opt_st_new, ps, grad) + end + end + ttime = time() - stime + # compute loss for test data + l = loss_gf_test(ps, test_data_dev...; is_testmode = true) + println() + @show round(ttime, digits=1), epoch, l.nLy, l.neg_log_prior, l.loss_penalty + # TODO log + end + res = nothing + ϕ = intϕ(ps) + #θP = !isempty(ϕ.ϕP) ? cpu_ca(apply_preserve_axes(transP, cpu_ca(ϕ).ϕq.μP)) : CA.ComponentVector{eltype(ϕ)}() + # TODO check which components live on gpu and which on cpu + ϕq_opt = ComponentArrayInterpreter(ϕq0)(gdevs.gdev_P(ϕ.ϕq)) + probo = HybridProblem(prob; ϕg=cpu_ca(ϕ).ϕg, ϕq=ϕq_opt) + (; ϕ, resopt=res, probo) +end + +function get_callback_epochs(epochs_callback; + n_site, n_batch, callback, loss_test) + n_per_epoch = n_site ÷ n_batch + if epochs_callback == 0 + callback_epochs = callback + else + callback_epochs = function(state, l) + if (state.iter == 1) || (state.iter % (n_per_epoch * epochs_callback) == 0) + l_test = loss_test(state)[1] + println("epoch = $((state.iter) ÷ n_per_epoch), iter = $(state.iter), loss_train=$l, loss_test=$l_test") + end + return callback(state, l) + end + end +end + +""" +A Solver that solves for the full posterior of a variational inference problem +([`AbstractHybridProblem`](@ref)). + +Parameters +- `alg`: The optimization algorithm to use for solving the problem, e.g. `Adam()` +- `n_MC`: The number of Monte Carlo samples to use for estimating the gradients of the loss function. +""" +struct HybridPosteriorSolverSite{A} <: AbstractHybridSolver + alg::A + n_MC::Int + n_MC_cap::Int +end +function HybridPosteriorSolverSite(; alg, n_MC=12, n_MC_cap=n_MC) + HybridPosteriorSolverSite(alg, n_MC, n_MC_cap) +end +function HybridPosteriorSolverSite(solver::HybridPosteriorSolverSite; + alg=solver.alg, + n_MC=solver.n_MC, + n_MC_cap=n_MC) + HybridPosteriorSolverSite(alg, n_MC, n_MC_cap) +end + +""" + solve(prob::AbstractHybridProblem, solver::HybridPosteriorSolverSite; epochs, ...) + +Perform the inversion of HVI Problem. + +Arguments +- `prob`: The AbstractHybridProblem to solve. +- `scenario`: Scenario to query prob, defaults to `Val(())`. +- `epochs`: number of epochs to train, i.e. number of passes through the whole dataset. + +Optional keyword arguments +- `rng`: Random generator, defaults to `Random.default_rng()`. +- `gdevs`: `NamedTuple` `(;gdev_M, gdev_P)` functions to move + computation and data of ML model on and PBM respectively + to gpu (e.g. `gpu_device()` or cpu (`identity`). + defaults to [`get_gdev_MP`](@ref)`(scenario)` +- `θmean_quant` default to `0.0`: deprecated +- `is_inferred`: set to `Val(true)` to activate type stability checks +- `is_omit_priors`: set to `Val(true)` to omit priors in the loss computation, which can be useful for debugging or if priors are not implemented for a specific scenario (e.g. on gpu) +- `clusters`: vector of cluster assignments for each site, defaults to each site being its own cluster. Clusters are used to compute the loss in a way that accounts for clustering of sites, which can be useful if there are many sites and the number of Monte Carlo samples is limited. +- `cluster_rep`: number of times to repeat each cluster in the loss computation, defaults to 1. Repeating clusters can be useful to effectively increase the number of Monte Carlo samples when the number of clusters is small. +- `epochs_callback`: number of epochs between progress output on evaluating testdata + +Returns a `NamedTuple` of +- `probo`: A copy of the HybridProblem, with updated optimized parameters +- `interpreters`: TODO +- `ϕ`: the optimized HVI parameters: a `ComponentVector` with entries + - `ϕg`: The ML model parameter vector, + - `ϕq`: `ComponentVector` of non-ML parameters, including + `μP`: `ComponentVector` of the mean global PBM parameters at unconstrained scale +- `θP`: `ComponentVector` of the mean global PBM parameters at constrained scale +- `resopt`: the structure returned by `Optimization.solve`. It can contain + more information on convergence. +""" +function CommonSolve.solve(prob::AbstractHybridProblem, solver::HybridPosteriorSolverSite; + epochs, + scenario::Val{scen}=Val(()), rng=Random.default_rng(), + gdevs = get_gdev_MP(scenario), + θmean_quant=0.0, + is_inferred::Val{is_infer} = Val(false), + is_omit_priors::Val{omit_priors} = Val(false), + approx = prob.approx, + clusters::AbstractVector{<:Integer} = + 1:first(get_hybridproblem_n_site_and_batch(prob; scenario)), + cluster_rep = 1, + epochs_callback = max(1, floor(Int, epochs/10)), + callback = (state, loss_val) -> false, + kwargs... +) where {scen, is_infer, omit_priors} + pt = get_hybridproblem_par_templates(prob; scenario) + cor_ends = get_hybridproblem_cor_ends(prob; scenario) + g, ϕg0 = get_hybridproblem_MLapplicator(prob; scenario) + (; transP, transM) = get_hybridproblem_transforms(prob; scenario) + pbm_covars = get_hybridproblem_pbmpar_covars(prob; scenario) + n_site, n_batch = get_hybridproblem_n_site_and_batch(prob; scenario) + n_sites_cluster = [count(==(element),clusters) for element in 1:maximum(clusters)] + frac_cluster_all = (1 / cluster_rep) ./ n_sites_cluster[clusters] + ranef_spec = get_hybridproblem_ranef(prob; scenario) + ranef = get_ranef_computer( + ranef_spec, keys(pt.θM), n_site, one(eltype(pt.θM))) + ϕq_ranef = setup_ϕq_ranef(ranef) + ϕq = CA.ComponentVector(get_hybridproblem_ϕq(prob; scenario), ranef = ϕq_ranef) + (; ϕ, interpreters) = init_hybrid_params(ϕg0, ϕq) + int_ϕq = interpreters.ϕq + int_ϕg_ϕq = interpreters.ϕg_ϕq + priors = get_hybridproblem_priors(prob; scenario) + priorsP = Tuple(priors[k] for k in keys(pt.θP)) + priorsM = Tuple(priors[k] for k in keys(pt.θM)) + zero_prior_logdensity = omit_priors ? 0f0 : get_zero_prior_logdensity( + priorsP, priorsM, pt.θP, pt.θM) + train_loader = get_hybridproblem_train_dataloader(prob; scenario) + test_data = get_hybridproblem_test_data(prob; scenario) + # i_test = rand(1:n_site, Integer(floor(n_site/10))) + # test_data = map(train_loader.data) do data_comp + # ndims(data_comp) == 2 ? data_comp[:, i_test] : data_comp[i_test] + # end + n_batch_test = size(test_data[1],2) + if first(train_loader)[1] isa CA.ComponentArray + @warn("ML model covariates (1) were provided as ComponentArray. " * + "Consider providing them as a plain array.") + end + if first(train_loader)[2] isa CA.ComponentArray + @warn("PBM drivers (2) were provided as ComponentArray. " * + "Consider providing them as a plain array.") + end + if gdevs.gdev_M isa MLDataDevices.AbstractGPUDevice + ϕ0_dev = gdevs.gdev_M(ϕ) + g_dev = gdevs.gdev_M(g) # zygote fails if gdev is a CPUDevice, although should be non-op + train_loader_dev = gdev_hybridproblem_dataloader(train_loader; gdevs) + test_data_dev = gdev_hybridproblem_data(test_data[keys(test_data)[1:5]]; gdevs) + else + ϕ0_dev = ϕ + g_dev = g + train_loader_dev = train_loader + test_data_dev = test_data[keys(test_data)[1:5]] + end + f = get_hybridproblem_PBmodel(prob; scenario) + f_test = create_nsite_applicator(f, n_batch_test) + if gdevs.gdev_P isa MLDataDevices.AbstractGPUDevice + f_dev = gdevs.gdev_P(f) #fmap(gdevs.gdev_P, f) + f_test_dev = gdevs.gdev_P(f_test) + else + f_dev = f + f_test_dev = f_test + end + + py = get_hybridproblem_neg_logden_obs(prob; scenario) + + penalty_computer = get_hybridproblem_penalty_computer(prob; scenario) + # intθP = ComponentArrayInterpreter(pt.θP) + # intθMs = ComponentArrayInterpreter((n_batch,), pt.θM) + + loss_elbo = get_loss_elbo( + g_dev, transP, transM, f_dev, py, n_batch; + n_MC = solver.n_MC, n_MC_cap = solver.n_MC_cap, cor_ends, + cdev=infer_cdev(gdevs), pbm_covars, + par_templates = pt, + #pt.θP, + int_ϕq, int_ϕg_ϕq, priorsP, priorsM, + is_omit_priors, zero_prior_logdensity, approx, penalty_computer, + ranef, + #intθMs, intθP, + frac_cluster_all, + ) + loss_elbo_test = get_loss_elbo( + g_dev, transP, transM, f_test_dev, py, n_batch_test; + solver.n_MC, solver.n_MC_cap, cor_ends, + cdev=infer_cdev(gdevs), pbm_covars, + par_templates = pt, + #pt.θP, + int_ϕq, int_ϕg_ϕq, priorsP, priorsM, + is_omit_priors, zero_prior_logdensity, approx, penalty_computer, + ranef, + #intθMs, intθP, + frac_cluster_all, + ) + + # test loss function once + # tmp = first(train_loader_dev) + # using ShareAdd + # @usingany Cthulhu + # @descend_code_warntype loss_elbo(ϕ0_dev, rng, first(train_loader_dev)...) + # omit for type stability in AD + @assert length(first(train_loader_dev)) == 5 + l0 = + #is_infer ? + # (Test.@inferred loss_elbo(ϕ0_dev, rng, first(train_loader_dev)...; is_testmode=true)) : + loss_elbo(ϕ0_dev, rng, first(train_loader_dev)...; is_testmode=false) + l0t = loss_elbo_test(ϕ0_dev, rng, test_data_dev...; is_testmode=true) + loss_test = let test_data_dev = test_data_dev + (state) -> loss_elbo_test(state.u, rng, test_data_dev...; is_testmode=true) + end + callback_epochs = get_callback_epochs(epochs_callback; + n_site, n_batch, callback, loss_test) + optf = Optimization.OptimizationFunction( + (ϕ, data) -> first(loss_elbo(ϕ, rng, data...; is_testmode=false)), + Optimization.AutoZygote()) + optprob = OptimizationProblem(optf, CA.getdata(ϕ0_dev), train_loader_dev) + res = Optimization.solve(optprob, solver.alg; callback = callback_epochs, epochs, kwargs...) + ϕc = interpreters.ϕg_ϕq(cpu_device()(res.u)) + ϕq = ϕc[Val(:ϕq)]; + ϕg = ϕc[Val(:ϕg)]; + probo = HybridProblem(prob; ϕg, ϕq) + θP = get_hybridproblem_θP(probo) + (; probo, interpreters, ϕ=ϕc, θP, resopt=res) +end + diff --git a/src/elbo_site.jl b/src/elbo_site.jl index a9a07be..c9afa56 100644 --- a/src/elbo_site.jl +++ b/src/elbo_site.jl @@ -232,7 +232,7 @@ function g_apply!(ϕm::AbstractMatrix{TF}, ϕg::AbstractVector{TG}, xM::Abstract g::AbstractModelApplicator, xMP::AbstractMatrix, is_testmode::Bool, - ϕms_mcs2D_buffer = nothing, # only required for 3D case + ϕms_mcs2D_buffer = nothing, # only required for pbm_covar_indices not nothing ) where {TG, TF} apply_model!(ϕm, g, xM, ϕg; is_testmode) # allocates view return nothing @@ -243,7 +243,7 @@ function g_apply!(ϕm::AbstractArray{TF,3}, ϕg::AbstractVector{TG}, xM::Abstrac ϕms_mcs2D_buffer = reshape(ϕm, size(ϕm,1), :) # allocates 48 bytes for an escaping view, preallocate ) where {TG, TF} update_xMP!(xMP, xM, ζsP, pbm_covar_indices) - #Main.@infiltrate_main + #stack the n_site * n_MC and copy back to array afterwards, if view points to other if pointer(parent(ϕms_mcs2D_buffer)) != pointer(parent(ϕm)) # if supplied a preallocated array (rather than using the default view) theń copy copyto!(ϕms_mcs2D_buffer, ϕm) # in order to call apply_model only once, stack n_sites * n_MC @@ -254,7 +254,6 @@ function g_apply!(ϕm::AbstractArray{TF,3}, ϕg::AbstractVector{TG}, xM::Abstrac end return nothing end - function update_xMP!(xMP::AbstractMatrix{TG}, xM::AbstractMatrix{TG}, ζsP::AbstractMatrix{TF}, pbm_covar_indices::AbstractVector{<:Number} ) where {TG, TF} @@ -265,6 +264,7 @@ function update_xMP!(xMP::AbstractMatrix{TG}, @inbounds for i in 1:n_site for j in 1:n_MC # Copy xM block + # for each site i there are n_MC columns, j, with the same xM ic = (i-1)*n_MC + j for k in 1:n_cov xMP[k, ic] = xM[k, i] @@ -277,6 +277,39 @@ function update_xMP!(xMP::AbstractMatrix{TG}, end end +function g_apply!(ϕm::AbstractMatrix{TF}, ϕg::AbstractVector{TG}, xM::AbstractMatrix{TG}, + ζP::AbstractVector{TF}, pbm_covar_indices::AbstractVector{<:Number}, + g::AbstractModelApplicator, + xMP::AbstractMatrix, + is_testmode::Bool, + ϕms_mcs2D_buffer = reshape(ϕm, size(ϕm,1), :) # allocates 48 bytes for an escaping view, preallocate + ) where {TG, TF} + update_xMP!(xMP, xM, ζP, pbm_covar_indices) # ζP vector version + @assert size(xMP,2) == size(ϕm,2) # same number of sites + apply_model!(ϕm, g, xMP, ϕg; is_testmode) + return nothing +end +function update_xMP!(xMP::AbstractMatrix{TG}, + xM::AbstractMatrix{TG}, ζP::AbstractVector{TF}, pbm_covar_indices::AbstractVector{<:Number} + ) where {TG, TF} + (n_θP,) = size(ζP) + n_cov, n_site = size(xM) + n_covP = length(pbm_covar_indices) + @assert size(xMP) == ((n_cov + n_covP) , n_site) + @inbounds for i in 1:n_site + # Copy xM block + for k in 1:n_cov + xMP[k, i] = xM[k, i] + end + # Fill pbm covariates block + for (k, idx) in enumerate(pbm_covar_indices) + xMP[n_cov + k, i] = (TG === TF) ? ζP[idx] : convert(TG, ζP[idx]) + end + end +end + + + function transformζ!(θs, trans, ζs::AbstractArray{TF}) where TF n_MC = size(ζs,2) # θs1 = exp.(ζs) diff --git a/src/pointloss_site.jl b/src/pointloss_site.jl new file mode 100644 index 0000000..0cac24e --- /dev/null +++ b/src/pointloss_site.jl @@ -0,0 +1,207 @@ +""" +elbo_helpers need to be initialized with new random numbers +in h.ζP and hi.ζM by calling randnPM! before. +By this way, we can compute the derivative corresponding to the forward pass +""" +function neg_pointloss_sites!( + pointloss_helpers::NamedTuple, # tuple of preallocated arrays + ϕ::NamedTuple, + sample_args::NamedTuple, + indiv_args::NamedTuple, + nljoint_args_fix::NamedTuple; +) + # for AD do not put it into closure + h = pointloss_helpers # preallocated μζP, dμζP, ζsP, ϕms, xMP, dxMP + (;ϕg, ϕqP, ϕqI) = ϕ + (;g, is_testmode, pbm_covar_indices, intϕqP, intϕqI, cor_ends, transP, transM) = sample_args + @assert keys(indiv_args)[1] == :xM + xM = indiv_args.xM; r2end = 2:length(indiv_args) + nljoint_args_inds = NamedTuple{keys(indiv_args)[r2end]}(values(indiv_args)[r2end]) + use_dc = h.helpers_sites[1].ζM isa PAT.DiffCache + (n_M,) = use_dc ? size(h.helpers_sites[1].ζM.du) : size(h.helpers_sites[1].ζM) + ϕqPc = intϕqP(ϕqP) + ϕqIc = intϕqI(ϕqI) + pbm_covar_indices = !isnothing(pbm_covar_indices) && isempty(pbm_covar_indices) ? nothing : pbm_covar_indices + # + h.ζP = CA.getdata(view(ϕqc,Val(:μζP))) + # (n_M x n_site) + # TODO g_apply without n_MC + g_apply!(h.ϕms, ϕg, xM, h.ζP, pbm_covar_indices, g, h.xMP, is_testmode) + exp_ladJacTP = transformζ!(h.θP, transP, h.ζP) # return value captures ladJacT + # so that one can provide its gradient to the pullback + ϕm_it = eachscol(h.ϕms) + template = ϕqI # only important for gradient + θP = h.θP + keys_nljoint_args_indi_val = Val(keys(nljoint_args_inds)) + # closure with approx, transM, and kwargs + 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 = add_names(nljoint_args_indi_tup, keys_nljoint_args_indi_val) + compute_pointlossi_z!(hi, nljoint_args_indi, ϕm, + ϕqIc, θP, cor_ends.M, transM, nljoint_args_fix) + end + #res_site = map(compute_pointlossi_z!, 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 + elbo_z = mapreduce(compute_nelboi_z_cl!, +, h.helpers_sites, rnormPM.M, zip_eachlastdims(nljoint_args_inds), ϕm_it) + # E = sum(x -> x.E, res_site) + # loglik = sum(x -> x.loglik, res_site) + # costTrans = sum(x -> x.costTrans, res_site) + #elbo = sum(first, res_site) - sum(h.logσ_ζP) + elbo = elbo_z - exp_ladJacTP - sum(h.logσ_ζP) +end + +add_names(tup, ::Val{KEYS}) where KEYS = NamedTuple{KEYS}(tup) + +function compute_pointlossi_z!(hi, + nljoint_args_indi, ϕm::AbstractArray, ϕqIc::AbstractArray, θP, + cor_endsM, transM, nljoint_args_fix, ) + # DM is 1 (vector) if no population covariates - constant across MC + # 2 (Matrix) if g is driven by population covariates - differs across MC + # on update -> sync corresponding function within grad_neg_elbo_sites + if hi.ζM isa PAT.DiffCache + # Gradiend called twice because of different dimensions need to pick the Dual + template = eltype(ϕqIc) <: ForwardDiff.Dual ? ϕqIc : θP + hi = map_leaves_nt(x -> PAT.get_tmp(x, template), hi) + end + n_θM = cor_endsM[end] + hi.ζM .= view(ϕm, 1:n_θM, :) + transformζ!(hi.θM, transM, hi.ζM) + exp_nL = exp_nLi!(hi.y, θP, hi.θM, nljoint_args_indi, nljoint_args_fix)[1] + pointlossi = exp_nL +end + +function prepare_pointloss_helpers(approx::AbstractHVIApproximation, + ϕg::AbstractArray{TG}, template_TF::AbstractArray{TF}; + n_θP, n_θM, n_site, n_MC, n_cov, n_covP, n_M, cor_ends, n_obs, + use_diff_cache::Val{use_dc} = Val(true), + diffchunk::ForwardDiff.Chunk{chunk} = ForwardDiff.Chunk(8), + ) where {TG, TF, use_dc, chunk} + his = Tuple((; + ζM = Vector{TF}(undef, n_θM), + θM = Vector{TF}(undef, n_θM), + y = Vector{TF}(undef, n_obs), + #sample_buffers = prepare_ind_sample_buffers(approx, cor_ends.M, template_TF), + ) for i in 1:n_site) + helpers_sites = use_dc ? map_leaves_nt(x -> PAT.DiffCache(x, chunk), his) : his + h = (; + ζP = Vector{TF}(undef, n_θP), + θP = Vector{TF}(undef, n_θP), + logσ_ζP = Vector{TF}(undef, n_θP), + ϕms = Matrix{TF}(undef, n_M, n_site), + ϕms_mcs2D_buffer = Matrix{TF}(undef, n_M, n_MC * n_site), + xMP = Matrix{TG}(undef, (n_cov + n_covP), n_MC * n_site), + sample_buffers = prepare_sample_buffers(approx, cor_ends.P, template_TF), + #diffchunk, + helpers_sites, + ) +end + +function check_pointloss_helpers(h::NamedTuple, xM::AbstractMatrix, pbm_covar_indices; + n_ϕg + ) + n_cov, n_site = size(xM) + n_covP = isnothing(pbm_covar_indices) ? 0 : length(pbm_covar_indices) + n_θP, n_MC = size(h.ζP) + n_M = size(h.ϕms, 1) + @assert size(h.ζP) == (n_θP, n_MC ) + @assert size(h.θP) == (n_θP, n_MC ) + #@assert size(h.dϕg) == (n_ϕg,) + @assert size(h.logσ_ζP) == (n_θP,) + @assert size(h.ϕms) == (n_M, n_site) + @assert size(h.ϕms_mcs) == (n_M, n_MC, n_site) + @assert size(h.xMP) == ((n_cov + n_covP), n_MC * n_site) + # + @assert length(h.helpers_sites) == n_site + hi = h.helpers_sites[1] + n_θM = size(hi.ζM.du, 1) + n_obs = size(hi.y.du, 1) + @assert size(hi.ζM.du) == (n_θM, n_MC) + @assert size(hi.θsM.du) == (n_θM, n_MC) + @assert size(hi.y.du) == (n_obs, n_MC) + @assert size(hi.logσ_ζM.du) == (n_θM,) +end + + +""" +compute the expected value of the negative log joint density of observations +and parameters +""" +function exp_nLi!( + y_pred::AbstractVector, + θP::AbstractVector, + θM::AbstractVector, + nljoint_args_indi::NamedTuple, + nljoint_args_fix::NamedTuple, + # f, py, + # priorsP, priorsM, + # penalty_computer = ZeroPenaltyComputer(), + # is_omit_priors, + # zero_prior_logdensity, + # intθP, intθMs, + # ranef::AbstractRandomEffectsComputer, + # frac_cluster_all, +) + (n_obs,) = size(y_pred) + (;xP, y_o, y_unc, i_indiv_train) = nljoint_args_indi + @assert size(y_o) == size(y_unc) == (n_obs,) + (;f!,fLogLik) = nljoint_args_fix + #θsPc = CA.ComponentMatrix(θsP, axθP, CA.FlatAxis()) + #θsMc = CA.ComponentMatrix(θsM, axθM, CA.FlatAxis()) + apply_model!(y_pred, f!, θP, θM, xP) + nL_reg = sum(y_pred) / n_MC + nL = fLogLik(y_pred, y_o, y_unc) / n_MC + # nL = (5 * sum(θsP) + 3 * sum(θsM)) / n_MC + (; nL,) + # loss_comps = neg_elbo_ζtf( + # ζsP_cpu[:,1:n_MC], ζsMs_tr_cpu[:,:,1:n_MC], σ, f, py, xP, y_ob, y_unc; + # n_MC_cap, transP, transMs, priorsP, priorsM, + # penalty_computer, ϕg, ϕqc, is_omit_priors, zero_prior_logdensity, + # itrain_sites, intθMs, intθP, ranef, frac_cluster_all) +end + +function fθpop_dummy1!(pred, θs, xPc::CA.ComponentVector{T}) where T + # variant that computes the model jointly across parameter vectors in a broadcast + # using @row_valviews + local n_θP, n_θM, n_MC, θsPc, θsMc + a1 = view(θs, Val(:a1), :) + @row_valviews a1,a2,a3,b1,b2,b3 = θs + # regression to nL = (5 * sum(θsP) + 3 * sum(θsM)) / n_MC + #n_θP = 3; n_θM = 3 + n_obs = length(view(xPc, Val(:S1))) + n_θ, n_MC = size(θs) + pred .= zero(eltype(pred)) + for θa in (a1,a2,a3) + pred .+= (xPc ./ xPc) .* T(5) .* θa' + end + for θb in (b1,b2,b3) + pred .+= (xPc ./ xPc) .* T(3) .* θb' + end + pred ./= n_obs + pred +end + +function fθpop_dummy2!(pred, θs, xPc::CA.ComponentVector{T}) where T + # variant that runs the model separately for each parameter vector + # using eachcol(θs) and @unpack_val + # using view(pred,:,j) to avoid intermediate allocations + pred .= zero(eltype(pred)) + n_obs = length(view(xPc, Val(:S1))) + for (j,θ) in enumerate(eachcol(θs)) + @unpack_val a1,a2,a3,b1,b2,b3 = θ # now scalar parameters + for θa in (a1,a2,a3) + view(pred,:,j) .+= (xPc ./ xPc) .* T(5) .* θa + end + for θb in (b1,b2,b3) + view(pred,:,j) .+= (xPc ./ xPc) .* T(3) .* θb + end + end + pred ./= n_obs + pred +end + +function fLogLik_dummy(y_pred, y_o, y_unc) + sum(y_pred) +end + diff --git a/test/test_elbo_site_doubleMM.jl b/test/test_elbo_site_doubleMM.jl index 1ba949b..a2cb272 100644 --- a/test/test_elbo_site_doubleMM.jl +++ b/test/test_elbo_site_doubleMM.jl @@ -240,7 +240,10 @@ pbm_covar_indices = SA.SVector{length(pbm_covars)}([findfirst(==(k), keys(par_te #_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) +end +@testset "HybridPointSolverSite" begin + end From 419d032626b6dadc9ee3fe342c57bfaa613650aa Mon Sep 17 00:00:00 2001 From: Thomas Wutzler Date: Wed, 21 Oct 2026 06:26:22 +0000 Subject: [PATCH 2/4] implement neg_pointloss_sites! --- src/HybridVariationalInference.jl | 3 + src/PBMApplicator.jl | 30 +- src/elbo_site.jl | 4 +- src/pointloss_site.jl | 147 ++------ src/pointloss_site_grad.jl | 594 ++++++++++++++++++++++++++++++ test/runtests.jl | 1 + test/test_elbo_site_doubleMM.jl | 127 ++----- 7 files changed, 672 insertions(+), 234 deletions(-) create mode 100644 src/pointloss_site_grad.jl diff --git a/src/HybridVariationalInference.jl b/src/HybridVariationalInference.jl index d54cfce..9f1d883 100644 --- a/src/HybridVariationalInference.jl +++ b/src/HybridVariationalInference.jl @@ -169,6 +169,9 @@ include("elbo_site_grad.jl") include("elbo_site_corr.jl") #export prepare_elbo_helpers +include("pointloss_site.jl") + + #export get_loss_ran_tr_f include("fit_ranef_site.jl") diff --git a/src/PBMApplicator.jl b/src/PBMApplicator.jl index e8b1ab4..91e5576 100644 --- a/src/PBMApplicator.jl +++ b/src/PBMApplicator.jl @@ -57,7 +57,9 @@ function create_nsite_applicator(app::PBMPopulationApplicator, n_MC; PBMPopulationApplicator(app.fθpop!, θFixm, θFixm_dc, app.axθcomp, app.ax_xP) end -function apply_model!(pred, app::PBMPopulationApplicator{AXS,AXP}, θsP::AbstractMatrix, θsM::AbstractMatrix, xP) where {AXS,AXP} +# single site, severl MC +function apply_model!(pred, app::PBMPopulationApplicator{AXS,AXP}, + θsP::AbstractMatrix, θsM::AbstractMatrix, xP::AbstractVector) where {AXS,AXP} # function tmpf(θsP, θsM, app) # @allocated Vcat(CA.getdata(θsP), CA.getdata(θsM), CA.getdata(app.θFixm)) # end @@ -75,8 +77,6 @@ function apply_model!(pred, app::PBMPopulationApplicator{AXS,AXP}, θsP::Abstrac # θc = CA.ComponentMatrix( # Vcat(CA.getdata(θsP), CA.getdata(θsM), CA.getdata(app.θFixm)), # AXC, CA.FlatAxis()) - # TODO provide PreallocationTools for app.θFixm - #θc = VcatCAs(AXS, CA.getdata(θsP), CA.getdata(θsM), CA.getdata(app.θFixm)) θFixm_dual = if eltype(θsP) <: ForwardDiff.Dual copyto!(PAT.get_tmp(app.θFixm_dc, CA.getdata(θsP)), app.θFixm) else @@ -91,7 +91,26 @@ function apply_model!(pred, app::PBMPopulationApplicator{AXS,AXP}, θsP::Abstrac pred end -function (app::PBMPopulationApplicator)(θP::AbstractMatrix{FT}, θMs_tr::AbstractArray, xP::AbstractMatrix) where FT +# single site, single MC +# causes allocations, better preallocate Matrix(n,1) +# function apply_model!(pred::AbstractVector, app::PBMPopulationApplicator{AXS,AXP}, θP::AbstractVector, θM::AbstractVector, xP::AbstractVector) where {AXS,AXP} +# local θc, xPc, θFixm_dual +# θFixm_dual = if eltype(θP) <: ForwardDiff.Dual +# copyto!(PAT.get_tmp(app.θFixm_dc, CA.getdata(θP)), app.θFixm) +# else +# app.θFixm +# end +# pred_mat = Base.ReshapedArray(pred, (length(pred), 1), ()) +# θP_mat = Base.ReshapedArray(pred, (length(θP), 1), ()) +# θM_mat = Base.ReshapedArray(pred, (length(θM), 1), ()) +# θc = VcatCAs(AXS, θP_mat, θM_mat, θFixm_dual) +# xPc = CA.ComponentVector(CA.getdata(xP), AXP) +# app.fθpop!(pred_mat, θc, xPc) +# pred +# end + +# several site, several MC +function apply_model(app::PBMPopulationApplicator,θP::AbstractMatrix{FT}, θMs_tr::AbstractArray, xP::AbstractMatrix) where FT local n_obs, n_site, n_par, n_MC, pred n_obs = get_n_obs(xP[:,1]) n_site, n_par, n_MC = size(θMs_tr) @@ -102,7 +121,8 @@ function (app::PBMPopulationApplicator)(θP::AbstractMatrix{FT}, θMs_tr::Abstra end (pred, nothing) end -function (app::PBMPopulationApplicator)(θP::AbstractVector{FT}, θMs_tr::AbstractMatrix, xP::AbstractMatrix) where FT +# several site, single MC +function apply_model(app::PBMPopulationApplicator, θP::AbstractVector{FT}, θMs_tr::AbstractMatrix, xP::AbstractMatrix) where FT # for a single MC-run local n_obs, n_site, pred, θsP, θsM n_obs = get_n_obs(xP[:,1]) diff --git a/src/elbo_site.jl b/src/elbo_site.jl index c9afa56..3a77312 100644 --- a/src/elbo_site.jl +++ b/src/elbo_site.jl @@ -48,7 +48,7 @@ function neg_elbo_sites!( # closure with approx, transM, and kwargs 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 = add_names(nljoint_args_indi_tup, keys_nljoint_args_indi_val) + nljoint_args_indi = name_tuple(nljoint_args_indi_tup, keys_nljoint_args_indi_val) compute_nelboi_z!(hi, approx, rnormM, nljoint_args_indi, ϕm, ϕqIc, θsP, cor_ends.M, transM, nljoint_args_fix) end @@ -64,7 +64,7 @@ function neg_elbo_sites!( elbo = elbo_z - exp_ladJacTP - sum(h.logσ_ζP) end -add_names(tup, ::Val{KEYS}) where KEYS = NamedTuple{KEYS}(tup) +name_tuple(tup, ::Val{KEYS}) where KEYS = NamedTuple{KEYS}(tup) diff --git a/src/pointloss_site.jl b/src/pointloss_site.jl index 0cac24e..259ad8c 100644 --- a/src/pointloss_site.jl +++ b/src/pointloss_site.jl @@ -23,20 +23,21 @@ function neg_pointloss_sites!( ϕqIc = intϕqI(ϕqI) pbm_covar_indices = !isnothing(pbm_covar_indices) && isempty(pbm_covar_indices) ? nothing : pbm_covar_indices # - h.ζP = CA.getdata(view(ϕqc,Val(:μζP))) + h.ζP .= CA.getdata(view(ϕqPc,Val(:μζP))) # (n_M x n_site) # TODO g_apply without n_MC - g_apply!(h.ϕms, ϕg, xM, h.ζP, pbm_covar_indices, g, h.xMP, is_testmode) + ϕms = isnothing(pbm_covar_indices) ? view(h.ϕms,:,:,1) : h.ϕms + g_apply!(ϕms, ϕg, xM, h.ζP, pbm_covar_indices, g, h.xMP, is_testmode) exp_ladJacTP = transformζ!(h.θP, transP, h.ζP) # return value captures ladJacT # so that one can provide its gradient to the pullback - ϕm_it = eachscol(h.ϕms) + ϕm_it = eachcol(view(h.ϕms,:,:,1)) template = ϕqI # only important for gradient θP = h.θP keys_nljoint_args_indi_val = Val(keys(nljoint_args_inds)) # closure with approx, transM, and kwargs - function compute_nelboi_z_cl!(hi, rnormM, nljoint_args_indi_tup, ϕm) + function compute_pointlossi_z_cl!(hi, nljoint_args_indi_tup, ϕm) #nljoint_args_indi = NamedTuple{keys_nljoint_args_indi}(nljoint_args_indi_tup) - nljoint_args_indi = add_names(nljoint_args_indi_tup, keys_nljoint_args_indi_val) + nljoint_args_indi = name_tuple(nljoint_args_indi_tup, keys_nljoint_args_indi_val) compute_pointlossi_z!(hi, nljoint_args_indi, ϕm, ϕqIc, θP, cor_ends.M, transM, nljoint_args_fix) end @@ -44,16 +45,12 @@ function neg_pointloss_sites!( #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 - elbo_z = mapreduce(compute_nelboi_z_cl!, +, h.helpers_sites, rnormPM.M, zip_eachlastdims(nljoint_args_inds), ϕm_it) - # E = sum(x -> x.E, res_site) - # loglik = sum(x -> x.loglik, res_site) - # costTrans = sum(x -> x.costTrans, res_site) - #elbo = sum(first, res_site) - sum(h.logσ_ζP) - elbo = elbo_z - exp_ladJacTP - sum(h.logσ_ζP) + #zind = Test.@inferred zip_eachlastdims(nljoint_args_inds) + zind = zip_eachlastdims(nljoint_args_inds) + pointloss_z = mapreduce(compute_pointlossi_z_cl!, +, h.helpers_sites, zind, ϕm_it) + pointloss = pointloss_z - exp_ladJacTP end -add_names(tup, ::Val{KEYS}) where KEYS = NamedTuple{KEYS}(tup) - function compute_pointlossi_z!(hi, nljoint_args_indi, ϕm::AbstractArray, ϕqIc::AbstractArray, θP, cor_endsM, transM, nljoint_args_fix, ) @@ -79,129 +76,37 @@ function prepare_pointloss_helpers(approx::AbstractHVIApproximation, diffchunk::ForwardDiff.Chunk{chunk} = ForwardDiff.Chunk(8), ) where {TG, TF, use_dc, chunk} his = Tuple((; - ζM = Vector{TF}(undef, n_θM), - θM = Vector{TF}(undef, n_θM), - y = Vector{TF}(undef, n_obs), - #sample_buffers = prepare_ind_sample_buffers(approx, cor_ends.M, template_TF), + ζM = Matrix{TF}(undef, n_θM, 1), + θM = Matrix{TF}(undef, n_θM, 1), + y = Matrix{TF}(undef, n_obs,1), ) for i in 1:n_site) helpers_sites = use_dc ? map_leaves_nt(x -> PAT.DiffCache(x, chunk), his) : his h = (; - ζP = Vector{TF}(undef, n_θP), - θP = Vector{TF}(undef, n_θP), - logσ_ζP = Vector{TF}(undef, n_θP), - ϕms = Matrix{TF}(undef, n_M, n_site), - ϕms_mcs2D_buffer = Matrix{TF}(undef, n_M, n_MC * n_site), - xMP = Matrix{TG}(undef, (n_cov + n_covP), n_MC * n_site), - sample_buffers = prepare_sample_buffers(approx, cor_ends.P, template_TF), - #diffchunk, + ζP = Matrix{TF}(undef, n_θP,1), + θP = Matrix{TF}(undef, n_θP,1), + ϕms = Array{TF}(undef, n_M, n_site,1), + xMP = Matrix{TG}(undef, (n_cov + n_covP), n_site), helpers_sites, ) end -function check_pointloss_helpers(h::NamedTuple, xM::AbstractMatrix, pbm_covar_indices; - n_ϕg - ) +function check_pointloss_helpers(h::NamedTuple, xM::AbstractMatrix, pbm_covar_indices) n_cov, n_site = size(xM) n_covP = isnothing(pbm_covar_indices) ? 0 : length(pbm_covar_indices) - n_θP, n_MC = size(h.ζP) + (n_θP,) = size(h.ζP) n_M = size(h.ϕms, 1) - @assert size(h.ζP) == (n_θP, n_MC ) - @assert size(h.θP) == (n_θP, n_MC ) - #@assert size(h.dϕg) == (n_ϕg,) - @assert size(h.logσ_ζP) == (n_θP,) - @assert size(h.ϕms) == (n_M, n_site) - @assert size(h.ϕms_mcs) == (n_M, n_MC, n_site) - @assert size(h.xMP) == ((n_cov + n_covP), n_MC * n_site) + @assert size(h.ζP) == (n_θP, 1) + @assert size(h.θP) == (n_θP, 1) + @assert size(h.ϕms) == (n_M, n_site,1) + @assert size(h.xMP) == ((n_cov + n_covP), n_site) # @assert length(h.helpers_sites) == n_site hi = h.helpers_sites[1] n_θM = size(hi.ζM.du, 1) n_obs = size(hi.y.du, 1) - @assert size(hi.ζM.du) == (n_θM, n_MC) - @assert size(hi.θsM.du) == (n_θM, n_MC) - @assert size(hi.y.du) == (n_obs, n_MC) - @assert size(hi.logσ_ζM.du) == (n_θM,) -end - - -""" -compute the expected value of the negative log joint density of observations -and parameters -""" -function exp_nLi!( - y_pred::AbstractVector, - θP::AbstractVector, - θM::AbstractVector, - nljoint_args_indi::NamedTuple, - nljoint_args_fix::NamedTuple, - # f, py, - # priorsP, priorsM, - # penalty_computer = ZeroPenaltyComputer(), - # is_omit_priors, - # zero_prior_logdensity, - # intθP, intθMs, - # ranef::AbstractRandomEffectsComputer, - # frac_cluster_all, -) - (n_obs,) = size(y_pred) - (;xP, y_o, y_unc, i_indiv_train) = nljoint_args_indi - @assert size(y_o) == size(y_unc) == (n_obs,) - (;f!,fLogLik) = nljoint_args_fix - #θsPc = CA.ComponentMatrix(θsP, axθP, CA.FlatAxis()) - #θsMc = CA.ComponentMatrix(θsM, axθM, CA.FlatAxis()) - apply_model!(y_pred, f!, θP, θM, xP) - nL_reg = sum(y_pred) / n_MC - nL = fLogLik(y_pred, y_o, y_unc) / n_MC - # nL = (5 * sum(θsP) + 3 * sum(θsM)) / n_MC - (; nL,) - # loss_comps = neg_elbo_ζtf( - # ζsP_cpu[:,1:n_MC], ζsMs_tr_cpu[:,:,1:n_MC], σ, f, py, xP, y_ob, y_unc; - # n_MC_cap, transP, transMs, priorsP, priorsM, - # penalty_computer, ϕg, ϕqc, is_omit_priors, zero_prior_logdensity, - # itrain_sites, intθMs, intθP, ranef, frac_cluster_all) + @assert size(hi.ζM.du) == (n_θM, 1) + @assert size(hi.θM.du) == (n_θM, 1) + @assert size(hi.y.du) == (n_obs, 1) end -function fθpop_dummy1!(pred, θs, xPc::CA.ComponentVector{T}) where T - # variant that computes the model jointly across parameter vectors in a broadcast - # using @row_valviews - local n_θP, n_θM, n_MC, θsPc, θsMc - a1 = view(θs, Val(:a1), :) - @row_valviews a1,a2,a3,b1,b2,b3 = θs - # regression to nL = (5 * sum(θsP) + 3 * sum(θsM)) / n_MC - #n_θP = 3; n_θM = 3 - n_obs = length(view(xPc, Val(:S1))) - n_θ, n_MC = size(θs) - pred .= zero(eltype(pred)) - for θa in (a1,a2,a3) - pred .+= (xPc ./ xPc) .* T(5) .* θa' - end - for θb in (b1,b2,b3) - pred .+= (xPc ./ xPc) .* T(3) .* θb' - end - pred ./= n_obs - pred -end - -function fθpop_dummy2!(pred, θs, xPc::CA.ComponentVector{T}) where T - # variant that runs the model separately for each parameter vector - # using eachcol(θs) and @unpack_val - # using view(pred,:,j) to avoid intermediate allocations - pred .= zero(eltype(pred)) - n_obs = length(view(xPc, Val(:S1))) - for (j,θ) in enumerate(eachcol(θs)) - @unpack_val a1,a2,a3,b1,b2,b3 = θ # now scalar parameters - for θa in (a1,a2,a3) - view(pred,:,j) .+= (xPc ./ xPc) .* T(5) .* θa - end - for θb in (b1,b2,b3) - view(pred,:,j) .+= (xPc ./ xPc) .* T(3) .* θb - end - end - pred ./= n_obs - pred -end - -function fLogLik_dummy(y_pred, y_o, y_unc) - sum(y_pred) -end diff --git a/src/pointloss_site_grad.jl b/src/pointloss_site_grad.jl new file mode 100644 index 0000000..29c80d2 --- /dev/null +++ b/src/pointloss_site_grad.jl @@ -0,0 +1,594 @@ + +function grad_neg_elbo_sites( + elbo_helpers::NamedTuple, # tuple of preallocated arrays + grad_elbo_helpers::NamedTuple, # tuple of preallocated arrays pullback closures + ϕ::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 + # approx::AbstractHVIApproximation, + # rnormPM::NamedTuple, # tuple of random numbers + # ϕg::AbstractVector{TG}, ϕqP::AbstractVector{TF}, ϕqI::AbstractVector{TF}, g, + # pbm_covar_indices::Union{Nothing,AbstractVector{<:Number}}, + # args...; + # nljoint_args_inds, # indices of sites in training set + # intϕqP, intϕqI, + # xM, + # cor_ends, + # transP::Stacked, transM::Stacked, + # is_testmode, + n_workers = 1, # TODO compute from executor and Distributed.nworkers and nthreads,... + executor::Transducers.Executor = Transducers.SequentialEx(), +) where NMC + h = elbo_helpers # preallocated μζP, dμζP, ζsP, ϕms, xMP, dxMP + (;ϕg, ϕqP, ϕqI) = ϕ + TF = eltype(ϕqP) + (;approx, g, is_testmode, pbm_covar_indices, intϕqP, intϕqI, cor_ends, transP, transM) = sample_args + @assert keys(indiv_args)[1] == :xM + xM = indiv_args.xM; r2end = 2:length(indiv_args) + nljoint_args_inds = NamedTuple{keys(indiv_args)[r2end]}(values(indiv_args)[r2end]) + use_ϕm_matrix = isnothing(pbm_covar_indices) + ϕqPc = intϕqP(ϕqP) + ϕqIc = intϕqI(ϕqI) + ϕm_buffer_key = use_ϕm_matrix ? :ϕms : :ϕms_mcs + h_ϕm = h[ϕm_buffer_key] + check_elbo_helpers(h, xM, pbm_covar_indices; n_ϕg = length(ϕg)) + n_cov, n_site = size(xM) + n_θP, n_MC = size(h.ζsP) + @assert n_MC == NMC + ax_θP = get_axis_θP(nljoint_args_fix.f!) + ax_θM = get_axis_θM(nljoint_args_fix.f!) + nljoint_args_indi = NamedTuple{keys(nljoint_args_inds)}( + first(zip_eachlastdims(nljoint_args_inds))) + gradh = !isempty(grad_elbo_helpers) ? grad_elbo_helpers : + #Test.@inferred prepare_gradelbo_helpers( + prepare_gradelbo_helpers( + ϕ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) + 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 + 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 + # + # parallel ForwardDiffGradient through forwarddiff_grad_nelboi_z! + cl = ForwardDiffGradNelboiZCl( + approx, ϕqIc, h.θsP, gradh.dϕmvecs, gradh.dθsP, gradh.hw_channel, cor_ends.M, transM, + Val(keys(nljoint_args_inds)), nljoint_args_fix, n_MC_val) + ϕm_it = eachslice(h_ϕm; dims = ndims(h_ϕm)) + init = (; + dϕqI = zero(TF) .* SA.SVector{axis_length(CA.getaxes(ϕqIc)[1])}(ϕqIc), + dθsP = zero(TF) .* SA.SVector{axis_length(ax_θP)*NMC}(h.θsP), ) + res = Folds.mapreduce(cl, make_tuple_reducer(+), + zip(h.helpers_sites, rnormPM.M, zip_eachlastdims(nljoint_args_inds), ϕm_it, 1:n_site), + executor; init) + # copy SharedArray into preallocated plain array for ∂elbo_∂θP + ∂elbo_∂ϕqm = gradh.∂elbo_∂ϕqm + copyto!(∂elbo_∂ϕqm, gradh.dϕmvecs) # copy from shared array + # copy StaticVector into preallocated ComponentVector + ∂elbo_∂θP = gradh.∂elbo_∂θP + 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) + ∂elbo_∂ladJacTP = -one(TF) + # + # pullback gradients of ϕqm -> gradh.dϕg and gradh.dζsP + gradh.pullback_g_apply!( + gradh.dϕg, gradh.∂elbo_∂ϕm_∂ζP, h_ϕm, ∂elbo_∂ϕqm, + ϕg, xM, h.ζsP, pbm_covar_indices, g, is_testmode) + # + # pullback gradients of ∂elbo_∂θP to ∂elbo_∂θP_∂ζP + ∂elbo_∂θP_∂ζP = gradh.∂elbo_∂θP_∂ζP + gradh.pullback_cl_transformζsP!( + ∂elbo_∂θP_∂ζP, + ∂elbo_∂θP, + ∂elbo_∂ladJacTP, + h.θsP, + transP, + h.ζsP, + ) + # + # pullback gradients of ∂elbo_∂ζP, ∂elbo_∂ϕm_∂ζP, and ∂elbo_∂logσ_ζP to dϕqP + gradh.pullback_cl_sample_ζsP!( + gradh.dϕqPc, + ∂elbo_∂θP_∂ζP + gradh.∂elbo_∂ϕm_∂ζP, gradh.∂elbo_∂logσ_ζP, + 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 + +""" +Callable to make deliver arguments that do not differ by individual to Foldl.mapreduce. +""" +struct ForwardDiffGradNelboiZCl{KEYS, TA, Tϕq, Tθ, TD, TP, THWC, TC, TM, TAF, NMC} + approx::TA + ϕqIc::Tϕq + θsP::Tθ + dϕmvecs::TD + dθsP::TP + hw_channel::THWC + corendsM::TC + transM::TM + # need to store KEYS in type parameter otherwise allocation in NamedTuple(keys) + keys_nljoint_args_indi::Val{KEYS} + nljoint_args_fix::TAF + n_MC_val::Val{NMC} +end +function (f::ForwardDiffGradNelboiZCl{KEYS})(tup) where KEYS + hi, rnormM, nljoint_args_indi_tup, ϕm, i = tup + #nljoint_args_indi = NamedTuple{f.keys_nljoint_args_indi}(nljoint_args_indi_tup) + nljoint_args_indi = NamedTuple{KEYS}(nljoint_args_indi_tup) + forwarddiff_grad_nelboi_z!(hi, f.approx, rnormM, nljoint_args_indi, ϕm, i, + f.ϕqIc, f.θsP, f.dϕmvecs, f.dθsP, f.hw_channel, f.corendsM, f.transM, + f.nljoint_args_fix, f.n_MC_val, + nothing, + #true, + ) +end + + +# ϕm vector version +# TOOD pass closure object rather than many parameters +function forwarddiff_grad_nelboi_z!(hi, approx::AbstractHVIApproximation, rnormM, + nljoint_args_indi, ϕm::AbstractVector{TF}, i, + ϕqIc, θsP, dϕmvecs, dθsP, hw_channel::Channel, cor_endsM, transM::Stacked, + nljoint_args_fix, n_MC_val::Val{NMC}, + omit_gradient=nothing) where {TF, NMC} + # need to calls to ForwardDiff.gradient! to avoid reshape + # group arrays of same dimensionality (here grad_ϕm and grad_ϕqI) into a Vcat + ax_ϕqI = CA.getaxes(ϕqIc)[1] + ax_ϕP = get_axis_θP(nljoint_args_fix.f!) + with_channel_element(hw_channel) do hwi_tup + hwi, hwi_dc = hwi_tup + #grad_conf = hwi.grad_conf + nelboi_z_mat, nelboi_z_vec = make_nelboiz_lazy_cls(hi, hwi_dc, approx, rnormM, nljoint_args_indi, + ϕm, ϕqIc, θsP, + cor_endsM, transM, nljoint_args_fix, + ax_ϕqI, Val(1), + ) + grad_θsP = if isnothing(omit_gradient) + ForwardDiff.gradient!(hwi.grad_θsP, nelboi_z_mat, CA.getdata(θsP), hwi.grad_conf_mat) + hwi.grad_θsP + else + θsP + end + grad_ϕm, grad_ϕqI = if isnothing(omit_gradient) + local ϕm_ϕqI + ϕm_ϕqI = Vcat(ϕm, CA.getdata(ϕqIc)) + ForwardDiff.gradient!(hwi.grad_ϕm_ϕqI, nelboi_z_vec, ϕm_ϕqI, hwi.grad_conf_vec) + hwi.grad_ϕm_ϕqI.args[1], hwi.grad_ϕm_ϕqI.args[2] + else + ϕm, CA.getdata(ϕqIc) + end + copyto!(view(dϕmvecs, :, i), grad_ϕm) + # returning SVector helps avoiding allocations during reduce + local res + res = (; + dϕqI = SA.SVector{axis_length(CA.getaxes(ϕqIc)[1])}(grad_ϕqI), + dθsP = SA.SVector{axis_length(ax_ϕP)*NMC}(grad_θsP), + ) + end +end + + +#@noinline CP.forwarddiff_grad_nelboi_z!(hi, approx, rnormM, nljoint_args_indi, ϕm_, i, +# ϕqIc_, θsP_, dϕmvecs_, dθsP_, hw_channel, cor_endsM, transM, nljoint_args_fix, omit_gradient) + + +# ϕm matrix version +function forwarddiff_grad_nelboi_z!(hi, approx::AbstractHVIApproximation, rnormM, + nljoint_args_indi, ϕm::AbstractMatrix{TF}, i, + ϕqIc, θsP, dϕmvecs, dθsP, hw_channel::Channel, cor_endsM, transM::Stacked, + nljoint_args_fix, n_MC_val::Val{NMC}, + omit_gradient=nothing) where {TF, NMC} + # need to calls to ForwardDiff.gradient! to avoid reshape + # group arrays of same dimensionality (here grad_ϕm and grad_ϕqI) into a Vcat + ax_ϕqI = CA.getaxes(ϕqIc)[1] + ax_ϕqP = get_axis_θP(nljoint_args_fix.f!) + with_channel_element(hw_channel) do hwi_tup + hwi, hwi_dc = hwi_tup + #grad_conf = hwi.grad_conf + nelboi_z_mat, nelboi_z_vec = make_nelboiz_lazy_cls(hi, hwi_dc, approx, rnormM, nljoint_args_indi, + ϕm, ϕqIc, θsP, + cor_endsM, transM, nljoint_args_fix, + ax_ϕqI, Val(2), ## 2-dimensional ϕm + ) + grad_θsP, grad_ϕm = if isnothing(omit_gradient) + local θsP_ϕm + θsP_ϕm = Vcat(θsP, ϕm) + ForwardDiff.gradient!(hwi.grad_θsP_ϕm, nelboi_z_mat, θsP_ϕm, hwi.grad_conf_mat) + hwi.grad_θsP_ϕm.args[1], hwi.grad_θsP_ϕm.args[2] + else + θsP, ϕm + end + grad_ϕqI = if isnothing(omit_gradient) + # TODO try reusing a view into grad_ϕm_ϕqI + ForwardDiff.gradient!(hwi.grad_ϕqI, nelboi_z_vec, CA.getdata(ϕqIc), hwi.grad_conf_vec) + hwi.grad_ϕqI + else + CA.getdata(ϕqIc) + end + copyto!(view(dϕmvecs, :, i), grad_ϕm) + # cannot use SVector for + # dθsP .+= grad_θsP # SharedArray, + # # returning SVector helps avoiding allocations during reduce + # local dϕqIc + # dϕqIc = SA.SVector{axis_length(CA.getaxes(ϕqIc)[1])}(grad_ϕqI) + local res + res = (; + dϕqI = SA.SVector{axis_length(CA.getaxes(ϕqIc)[1])}(grad_ϕqI), + dθsP = SA.SVector{axis_length(ax_ϕqP)*NMC}(grad_θsP), + ) + end +end + + +function make_nelboiz_lazy_cls(hi, hwi_dc, approx::AbstractHVIApproximation, rnormM, nljoint_args_indi, + ϕm, ϕqIc, θsP, + cor_endsM, transM, nljoint_args_fix, + ax_ϕqI::CA.Axis, ndim_ϕ::Val{1}, + ) + # mat, here holds only θsP + # vec, is a Vcat(ϕqI, ϕm) + # no views involved here + function nelboiz_cl_mat(mat::AbstractMatrix) + local θsP_dual, ϕm_dual, ϕqIc_dual + θsP_dual = mat + ϕm_dual = copyto!(PAT.get_tmp(hwi_dc.ϕm_dc, mat), ϕm) + ϕqIc_dual = CA.ComponentArray( + copyto!(PAT.get_tmp(hwi_dc.ϕqI_dc, mat), ϕqIc), CA.getaxes(ϕqIc)) + compute_nelboi_z!( + hi, approx, rnormM, nljoint_args_indi, + ϕm_dual, ϕqIc_dual, θsP_dual, + cor_endsM, transM, nljoint_args_fix, + )[1] + end + function nelboiz_cl_vec(vec::Vcat) + local θsP_dual, ϕm_dual, ϕqIc_dual + ϕm_dual = vec.args[1] + ϕqIc_dual = CA.ComponentVector(vec.args[2], ax_ϕqI) + θsP_dual = copyto!(PAT.get_tmp(hwi_dc.θsP_dc, vec), θsP) # reshape(reinterpred( + compute_nelboi_z!( + hi, approx, rnormM, nljoint_args_indi, + ϕm_dual, ϕqIc_dual, θsP_dual, + cor_endsM, transM, nljoint_args_fix, + )[1] + end + nelboiz_cl_mat, nelboiz_cl_vec +end + +# ϕm Matrix case +function make_nelboiz_lazy_cls(hi, hwi_dc, approx::AbstractHVIApproximation, rnormM, nljoint_args_indi, + ϕm, ϕqIc, θsP, + cor_endsM, transM, nljoint_args_fix, + ax_ϕqI::CA.Axis, ndim_ϕ::Val{2}, + ) + # mat is Vcat(θsP, ϕm) + # vec is ϕqI + function nelboiz_cl_mat(mat::Vcat) + local θsP_dual, ϕm_dual, ϕqIc_dual + θsP_dual = mat.args[1] + ϕm_dual = mat.args[2] + ϕqIc_dual = CA.ComponentArray( + copyto!(PAT.get_tmp(hwi_dc.ϕqI_dc, mat), ϕqIc), CA.getaxes(ϕqIc)) + compute_nelboi_z!( + hi, approx, rnormM, nljoint_args_indi, + ϕm_dual, ϕqIc_dual, θsP_dual, + cor_endsM, transM, nljoint_args_fix, + )[1] + end + function nelboiz_cl_vec(vec::AbstractVector) + local θsP_dual, ϕm_dual, ϕqIc_dual + ϕqIc_dual = CA.ComponentVector(vec, ax_ϕqI) + ϕm_dual = copyto!(PAT.get_tmp(hwi_dc.ϕm_dc, vec), ϕm) + θsP_dual = copyto!(PAT.get_tmp(hwi_dc.θsP_dc, vec), θsP) + compute_nelboi_z!( + hi, approx, rnormM, nljoint_args_indi, + ϕm_dual, ϕqIc_dual, θsP_dual, + cor_endsM, transM, nljoint_args_fix, + )[1] + end + nelboiz_cl_mat, nelboiz_cl_vec +end + + +function get_pullback_cl_sample_ζsP(::AbstractArray{TF}; n_θP, n_MC, sample_buffers) where TF + #dϕqc, dζsP, dlogσ_ζP, ζsP, logσ_ζP, rnormP, ϕqc) + ζsP_ = Matrix{TF}(undef, n_θP, n_MC) + logσ_ζP_ = Vector{TF}(undef, n_θP) + dζsP_ = similar(ζsP_) # allocate space for derivatives + dlogσ_ζP_ = similar(logσ_ζP_) + drnormP = similar(ζsP_) + dsample_buffers = map_leaves_nt(similar, sample_buffers) # allocate space for derivatives + # + function pullback_cl_sample_ζsP!(dϕqc, dζsP, dlogσ_ζP, ζsP, logσ_ζP, + approx::AbstractHVIApproximation, rnormP, ϕqc, cor_endsP, sample_buffers::NamedTuple) + Enzyme.make_zero!(dϕqc) + Enzyme.make_zero!(dsample_buffers) + fill!(dϕqc, 0) # the derivative to compute + copyto!(dζsP_, dζsP) # input cotangents (modified in-place by Enzyme) + copyto!(dlogσ_ζP_, dlogσ_ζP) + copyto!(ζsP_, ζsP) # modified in-place by Enzyme + copyto!(logσ_ζP_, logσ_ζP) + EnzBuffers = isempty(sample_buffers) ? Enzyme.Const(sample_buffers) : + Enzyme.Duplicated(sample_buffers, dsample_buffers) + Enzyme.autodiff( + Enzyme.Reverse, + sample_ζsP!, + Enzyme.Duplicated(ζsP_, dζsP_), + Enzyme.Duplicated(logσ_ζP_, dlogσ_ζP_), + Enzyme.Const(approx), + Enzyme.DuplicatedNoNeed(rnormP, drnormP), + Enzyme.Duplicated(ϕqc, dϕqc), + Enzyme.Const(cor_endsP), + EnzBuffers, + ) + end +end + +function get_pullback_g_apply(::AbstractArray{TG}, ::AbstractArray{TF}; + n_θP, n_cov, n_covP, n_MC, n_site, n_M, + ) where {TG, TF} + xMP_ = Matrix{TG}(undef, (n_cov + n_covP), n_MC * n_site) + dxMP_ = similar(xMP_) + ϕms_ = Matrix{TF}(undef, n_M, n_site) + ϕms_mcs_ = Array{TF,3}(undef, n_M, n_MC, n_site) + dϕms_ = similar(ϕms_) + dϕms_mcs_ = similar(ϕms_mcs_) + #dζsP_ = Matrix{TF}(undef, n_θP, n_MC) + # + function pullback_g_apply!(dϕg, dζsP, ϕms, dϕm, ϕg, xM, ζsP, + pbm_covar_indices, g, is_testmode) + ϕms_buffer = isnothing(pbm_covar_indices) ? ϕms_ : ϕms_mcs_ + dϕms_buffer = isnothing(pbm_covar_indices) ? dϕms_ : dϕms_mcs_ + # assert that buffers were constructed with correct sizes + n_cov_f, n_site_f = size(xM) + n_covP_f = isnothing(pbm_covar_indices) ? 0 : length(pbm_covar_indices) + n_MC_f = size(ζsP,2) + @assert (n_cov, n_covP, n_MC, n_site) == (n_cov_f, n_covP_f, n_MC_f, n_site_f) + @assert size(dϕms_buffer) == size(dϕm) + + fill!(dϕg, zero(eltype(dϕg))) + fill!(dζsP, zero(eltype(dζsP))) # also output cotangent + fill!(dxMP_, zero(eltype(dxMP_))) + # primal will be updated as in the forward, but shadow needs to be preserved + copyto!(ϕms_buffer, ϕms) # copy to avoid modifying ϕm (although should be the same) + copyto!(dϕms_buffer, dϕm) # copy to avoid modifying dϕm + #copyto!(dζsP_, dζsP) # output, does not need to be preserved + # + Enzyme.autodiff( + Enzyme.Reverse, + #does not make SimpleChains work - currently cannot use SimpleChains + #Enzyme.set_runtime_activity(Enzyme.Reverse), # TODO only activate for SimpleChains + g_apply!, + Enzyme.Duplicated(ϕms_buffer, dϕms_buffer), + Enzyme.Duplicated(ϕg, dϕg), + Enzyme.Const(xM), + Enzyme.Duplicated(ζsP, dζsP), + Enzyme.Const(pbm_covar_indices), + Enzyme.Const(g), + Enzyme.Duplicated(xMP_, dxMP_), + Enzyme.Const(is_testmode), + ) + end +end + +# two return values: primal and derivative +# given coderiv dy, parameters and helpers +# dϕm -> dϕg, dζsP +function get_pullback_g_apply_ϕms_mcs2D_buffer(::AbstractArray{TG}, ::AbstractArray{TF}; + n_θP, n_cov, n_covP, n_MC, n_site, n_M, + ) where {TG, TF} + xMP_ = Matrix{TG}(undef, (n_cov + n_covP), n_MC * n_site) + dxMP_ = similar(xMP_) + ϕms_ = Matrix{TF}(undef, n_M, n_site) + ϕms_mcs_ = Array{TF,3}(undef, n_M, n_MC, n_site) + dϕms_ = similar(ϕms_) + dϕms_mcs_ = similar(ϕms_mcs_) + #dζsP_ = Matrix{TF}(undef, n_θP, n_MC) + ϕms_mcs2D_buffer_ = Matrix{TF}(undef, n_M, n_MC * n_site) + dϕms_mcs2D_buffer_ = similar(ϕms_mcs2D_buffer_) + # + function pullback_g_apply_ϕms_mcs2D_buffer!(dϕg, dζsP, ϕms, dϕm, ϕg, xM, ζsP, + pbm_covar_indices, g, is_testmode, ϕms_mcs2D_buffer) + ϕms_buffer = isnothing(pbm_covar_indices) ? ϕms_ : ϕms_mcs_ + dϕms_buffer = isnothing(pbm_covar_indices) ? dϕms_ : dϕms_mcs_ + # assert that buffers were constructed with correct sizes + n_cov_f, n_site_f = size(xM) + n_covP_f = isnothing(pbm_covar_indices) ? 0 : length(pbm_covar_indices) + n_MC_f = size(ζsP,1) + @assert (n_cov, n_covP, n_MC, n_site) == (n_cov_f, n_covP_f, n_MC_f, n_site_f) + @assert size(dϕms_buffer) == size(dϕm) + @assert size(ϕms_mcs2D_buffer_) == size(ϕms_mcs2D_buffer) + + fill!(dϕg, zero(eltype(dϕg))) + fill!(dζsP, zero(eltype(dζsP))) # also output cotangent + fill!(dxMP_, zero(eltype(dxMP_))) + fill!(dϕms_mcs2D_buffer_, zero(eltype(dϕms_mcs2D_buffer_))) + # primal will be updated as in the forward, but shadow needs to be preserved + copyto!(ϕms_buffer, ϕms) # copy to avoid modifying ϕm (although should be the same) + copyto!(dϕms_buffer, dϕm) # copy to avoid modifying dϕm + #copyto!(dζsP_, dζsP) # output, does not need to be preserved + copyto!(ϕms_mcs2D_buffer_, ϕms_mcs2D_buffer) # copy to avoid modifying dϕm + # + Enzyme.autodiff( + Enzyme.Reverse, g_apply!, + Enzyme.Duplicated(ϕms_buffer, dϕms_buffer), + Enzyme.Duplicated(ϕg, dϕg), + Enzyme.Const(xM), + Enzyme.Duplicated(ζsP, dζsP), + Enzyme.Const(pbm_covar_indices), + Enzyme.Const(g), + Enzyme.Duplicated(xMP_, dxMP_), + Enzyme.Const(is_testmode), + Enzyme.Duplicated(ϕms_mcs2D_buffer_, dϕms_mcs2D_buffer_), + ) + end +end + +function get_pullback_cl_transformζ!(::AbstractArray{TF}; n_θ, n_MC) where {TF} + θs_buffer = Matrix{TF}(undef, n_θ, n_MC) + dθs_buffer = similar(θs_buffer) + # Enzyme seeds an Active return value with one, so to seed its cotangent + # with an arbitrary dladJacTP we fold it in as a scalar factor, relying on + # linearity of the reverse-mode adjoint: the pullback then delivers + # dladJacTP * ∂ladJacT/∂(·) to θs and ζs, as with the former Ref seeding. + function pullback_cl_transformζ!(dζs, dθs, dladJacTP, θs, trans, ζs) + fill!(dζs, zero(eltype(dζs))) + copyto!(θs_buffer, θs) + copyto!(dθs_buffer, dθs) + # Trick of seeding the active return value different to unity: + # Here dladJacTP is captured from the closure (an Active-compatible scalar), + # the returned Active value's unit seed gets multiplied by dladJacTP, + # and the adjoints of θs/ζs come out as dladJacTP * ∂/∂(·) + Enzyme.autodiff( + Enzyme.Reverse, + (θs_, trans_, ζs_) -> dladJacTP * transformζ!(θs_, trans_, ζs_), + Enzyme.Duplicated(θs_buffer, dθs_buffer), + Enzyme.Const(trans), + Enzyme.Duplicated(ζs, dζs), + ) + end +end + +function prepare_gradelbo_helpers( + ϕqIc, ϕm::AbstractArray{TF,ND}, θsP, + ϕg::AbstractVector{TG}, ϕqPc::AbstractVector{TF}, + approx::AbstractHVIApproximation; + pbm_covar_indices, n_workers, + h, rnormMi, nljoint_args_indi, + n_site, n_cov, cor_ends, + transP::Stacked, transM::Stacked, + nljoint_args_fix, + n_MC_val::Val{NMC}, + diffchunk::ForwardDiff.Chunk{chunk} = ForwardDiff.Chunk(8), + ) where {TG, TF, ND, NMC, chunk} + hi1 = h.helpers_sites[1] + ax_θP = get_axis_θP(nljoint_args_fix.f!) + ax_θM = get_axis_θM(nljoint_args_fix.f!) + # TODO get n_site, n_cov from + cv_grad = CA.ComponentArray(; + ϕqIc, + ϕm, # = selectdim(ϕms, ndims(ϕm), 1), + θsP, + ) + # ϕqIc = cv_grad[Val(:ϕqIc)] + # ϕm = cv_grad[Val(:ϕm)] + # θsP = cv_grad[Val(:θsP)] + use_ϕm_vector = isnothing(pbm_covar_indices) # for one site, provide singe ϕm across MC + n_covP = use_ϕm_vector ? 0 : length(pbm_covar_indices) + n_ϕmvec = length(ϕm) #use_ϕm_matrix ? n_M : n_M * n_MC + n_θP, n_MC = size(θsP) + @assert n_MC == NMC + @assert n_θP == axis_length(ax_θP) + n_M = size(ϕm,1) + # To sync across procs/threads, use a Channel + # https://juliafolds2.github.io/OhMyThreads.jl/stable/literate/tls/tls/#The-safe-way:-Channel + # + @assert length(nljoint_args_indi.xP) > 1 + # nelboi_z = make_nelboiz_cl(hi1, approx, rnormMi, nljoint_args_indi, CA.getaxes(cv_grad), + # cor_ends.M, transM, nljoint_args_fix) + hw_dcs_raw = (; + ϕm_dc = similar(ϕm), + # callocation on using ComponentArray (_restructure), need to use plain vector + ϕqI_dc = similar(CA.getdata(ϕqIc)), + θsP_dc = similar(θsP), + ) + # make sure, each worker gets its own copy, use similar + hw_dcs = (hw_dcs_raw) -> map_leaves_nt(x -> PAT.DiffCache(similar(x), chunk), hw_dcs_raw) + nelboi_z_mat, nelboi_z_vec = make_nelboiz_lazy_cls(hi1, hw_dcs(hw_dcs_raw), approx, rnormMi, nljoint_args_indi, + ϕm, ϕqIc, θsP, + cor_ends.M, transM, nljoint_args_fix, + CA.getaxes(ϕqIc)[1], Val(ND), + ) + grad_conf_mat, grad_conf_vec = if use_ϕm_vector + local grad_ϕm_ϕqI1, diffchunks_vec, diffchunks_mat + diffchunk_vec = ForwardDiff.Chunk(min(chunk, + axis_length(CA.getaxes(ϕqIc)[1]) + axis_length(ax_θM))) + diffchunk_mat = ForwardDiff.Chunk(min(chunk, + axis_length(CA.getaxes(ϕqPc)[1]))) + grad_conf_mat_ϕmvec() = ForwardDiff.GradientConfig(nelboi_z_mat, CA.getdata(θsP), diffchunk_mat) + grad_ϕm_ϕqI1 = grad_ϕm_ϕqI(ϕm, ϕqIc) + function grad_conf_vec_ϕmvec() + cfg_vec = ForwardDiff.GradientConfig(nelboi_z_vec, grad_ϕm_ϕqI1, diffchunk_vec) + update_duals_Vcat(cfg_vec, grad_ϕm_ϕqI1) + end + grad_conf_mat_ϕmvec, grad_conf_vec_ϕmvec + else + local rad_θP_ϕm1, diffchunks_vec, diffchunks_mat + diffchunks_vec = ForwardDiff.Chunk(min(chunk, + axis_length(CA.getaxes(ϕqIc)[1]))) + diffchunks_mat = ForwardDiff.Chunk(min(chunk, + NMC*(axis_length(CA.getaxes(ϕqPc)[1]) + axis_length(ax_θM)))) + grad_conf_vec_ϕmmat() = ForwardDiff.GradientConfig(nelboi_z_vec, CA.getdata(ϕqIc), diffchunks_vec) + rad_θP_ϕm1 = grad_θsP_ϕm(ϕm, θsP) + function grad_conf_mat_ϕmmat() + cfg_mat = ForwardDiff.GradientConfig(nelboi_z_mat, rad_θP_ϕm1, diffchunks_mat) + update_duals_Vcat(cfg_mat, rad_θP_ϕm1) + end + grad_conf_mat_ϕmmat, grad_conf_vec_ϕmmat + end + get_helpers_worker = () -> begin + (; + grad_conf_vec = grad_conf_vec(), + grad_conf_mat = grad_conf_mat(), + grad_θsP = similar(θsP), + grad_ϕm_ϕqI = grad_ϕm_ϕqI(ϕm, ϕqIc), + grad_θsP_ϕm= grad_θsP_ϕm(ϕm, θsP), + grad_ϕqI = similar(CA.getdata(ϕqIc)), + ) + end + h1 = get_helpers_worker() + h1_dc = hw_dcs(hw_dcs_raw) + h1_tup = (h1, h1_dc) + # put tuples of h1 an d1_dc into the channel + hw_channel = Channel{typeof(h1_tup)}(n_workers) # workers + parallel threads + put!(hw_channel, h1_tup) + foreach(2:n_workers) do _ + put!(hw_channel, (get_helpers_worker(), hw_dcs(hw_dcs_raw))) + end + (; + dϕg = similar(ϕg), + dϕqPc = similar(ϕqPc), + dϕqIc = similar(ϕqIc), + dθsP = similar(θsP), + ∂elbo_∂θP = similar(θsP), + ∂elbo_∂ϕm_∂ζP = similar(θsP), + ∂elbo_∂logσ_ζP = Vector{TF}(undef, size(θsP,1)), + ∂elbo_∂θP_∂ζP = Matrix{TF}(undef, size(θsP)), + ∂elbo_∂ϕqm = Array{TF}(undef, size(ϕm)..., n_site), + hw_channel, + pullback_cl_sample_ζsP! = get_pullback_cl_sample_ζsP(CA.getdata(ϕqPc); + n_θP, n_MC, sample_buffers = h.sample_buffers), + pullback_cl_transformζsP! = get_pullback_cl_transformζ!(CA.getdata(ϕqPc); n_θ = n_θP, n_MC), + pullback_g_apply! = get_pullback_g_apply( + ϕg, ϕqPc; n_θP, n_cov, n_covP, n_MC, n_site, n_M), + # SharedArrays last because of problems with ReferenceRevision + dϕmvecs = SharedArrays.SharedArray{TF}(n_ϕmvec, n_site), + #dθsP = SharedArrays.SharedArray{TF}(n_θP, n_MC), + ) +end +@inline grad_ϕm_ϕqI(ϕm::AbstractVector, ϕqIc) = Vcat(similar(ϕm), similar(CA.getdata(ϕqIc))) +@inline grad_ϕm_ϕqI(ϕm::AbstractMatrix, ϕqIc) = Vcat(similar(ϕm[1:0,:1]), similar(CA.getdata(ϕqIc))) # not used but preallocated +@inline grad_θsP_ϕm(ϕm::AbstractVector, θsP) = Vcat(similar(θsP[:,1:0]), similar(θsP[:,1:0])) # not used but preallocated +@inline grad_θsP_ϕm(ϕm::AbstractMatrix, θsP) = Vcat(similar(θsP), similar(ϕm)) + + +@inline function update_duals_Vcat(cfg::ForwardDiff.GradientConfig{T,V,N}, vc::Vcat) where {T,V,N} + # similar Vcat returns a vector, need to convert dual of cfg to Vcat + # no view, but copy with index_at_dim + duals_vcat = Vcat(index_at_dim(cfg.duals, axes(vc.args[1], 1); dim=1), + index_at_dim(cfg.duals, axes(vc.args[2], 1); dim=1)); + cfg = ForwardDiff.GradientConfig{T,V,N,typeof(duals_vcat)}(cfg.seeds, duals_vcat) + @assert cfg.duals isa Vcat + cfg +end + + + + diff --git a/test/runtests.jl b/test/runtests.jl index 4b22c30..b1bc06a 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -48,6 +48,7 @@ const GROUP = get(ENV, "GROUP", "All") # defined in in CI.yml @safetestset "BMApplicatorMC" include("test/test_PBMApplicatorMC.jl") @safetestset "util" include("test/test_util.jl") @safetestset "elbo_sites" include("test/test_elbo_site.jl") + @safetestset "elbo_sites_doubleMM" include("test/test_elbo_site_doubleMM.jl") end end diff --git a/test/test_elbo_site_doubleMM.jl b/test/test_elbo_site_doubleMM.jl index a2cb272..34a62c4 100644 --- a/test/test_elbo_site_doubleMM.jl +++ b/test/test_elbo_site_doubleMM.jl @@ -87,108 +87,6 @@ pbm_covar_indices = SA.SVector{length(pbm_covars)}([findfirst(==(k), keys(par_te rnormPM = CP.prepare_rnorm(ϕqP; n_θP, n_θM, n_MC, n_site) #------------------------------------------------------------- - -# @testset "compute_nelboi_z!" begin -# ϕqIc = intϕqI(ϕqI) -# ϕqPc = intϕqP(ϕqP) -# @assert keys(indiv_args1)[1] == :xM -# 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 -# CP.g_apply!(h0.ϕms, ϕg, xM, h0.ζsP, nothing, g, h0.xMP, false) -# hi1 = h0.helpers_sites[1] -# nljoint_args_ind1 = NamedTuple{keys(nljoint_args_inds)}( -# first(zip_eachlastdims(nljoint_args_inds))) -# rnormM1 = rnormPM.M[1] -# ϕms1 = h0.ϕms[:,1] -# θsP1 = h0.θsP -# CP.compute_nelboi_z!(hi1, 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, -# ϕ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) -# # -# # 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, -# 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) -# for _ in 1:n -# @noinline CP.compute_nelboi_z!(h21, 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) -# # 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 -# end -# alloc_compute_nelboi_z(h21, 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) -# 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 -# # measured region so that their construction is not counted by @allocated. -# function loop_forwarddiff_grad_nelboi_z( -# n, hi, approx, rnormM, nljoint_args_indi, i, ϕm_, ϕqIc_, θsP_, dϕmvecs_, dθsP_, hw_channel::Channel, omit_gradient, cor_endsM, transM::Stacked, nljoint_args_fix, n_MC_val) -# for _ in 1:n -# @noinline CP.forwarddiff_grad_nelboi_z!(hi, approx, rnormM, nljoint_args_indi, ϕm_, i, -# ϕqIc_, θsP_, dϕmvecs_, dθsP_, hw_channel, cor_endsM, transM, nljoint_args_fix, n_MC_val, omit_gradient) -# end -# nothing -# end -# ϕm_ = view(inputs, Val(:ϕm)) -# ϕqIc_ = view(inputs, Val(:ϕqIc)) -# θsP_ = view(inputs, Val(:θsP)) -# dϕmvecs = gradh2.dϕmvecs -# dθsP = gradh2.dθsP -# function alloc_forwarddiff_grad_nelboi_z(h21, approx, rnormM1, nljoint_args_indi, ϕm_, ϕqIc_, θsP_, dϕmvecs, dθsP_, hw_channel::Channel, cor_endsM, transM::Stacked, nljoint_args_fix, n_MC_val) -# loop_forwarddiff_grad_nelboi_z(1, h21, approx, rnormM1, nljoint_args_indi, 1, ϕm_, ϕqIc_, θsP_, dϕmvecs, dθsP_, hw_channel, true, cor_endsM, transM, nljoint_args_fix, n_MC_val) -# @test (@allocated loop_forwarddiff_grad_nelboi_z(100, h21, approx, rnormM1, nljoint_args_indi, 1, ϕm_, ϕqIc_, θsP_, dϕmvecs, dθsP_, hw_channel, true, cor_endsM, transM, nljoint_args_fix, n_MC_val)) == 0 -# loop_forwarddiff_grad_nelboi_z(1, h21, approx, rnormM1, nljoint_args_indi, 1, ϕm_, ϕqIc_, θsP_, dϕmvecs, dθsP_, hw_channel, nothing, cor_endsM, transM, nljoint_args_fix, n_MC_val) -# @test (@allocated loop_forwarddiff_grad_nelboi_z(100, h21, approx, rnormM1, nljoint_args_indi, 1, ϕm_, ϕqIc_, θsP_, dϕmvecs, dθsP_, hw_channel, nothing, cor_endsM, transM, nljoint_args_fix, n_MC_val)) == 0 -# end -# alloc_forwarddiff_grad_nelboi_z(h21, approx, rnormM1, nljoint_args_ind1, ϕm_, ϕqIc_, θsP_, dϕmvecs, dθsP, hw_channel, cor_ends.M, transM, nljoint_args_fix, n_MC_val) - -# #@usingany Cthulhu -# #@descend_code_warntype tmp_g(h21, rnormM1, nljoint_args_ind1, inputs, 1, dϕmvecs, nothing, hw_channel) -# #@descend_code_warntype tmp_g2(h21, rnormM1, nljoint_args_ind1, inputs, 1, dϕmvecs, nothing, hw_channel) -# #@usingany BenchmarkTools -# #@benchmark tmp_g($h21, $rnormM1, $nljoint_args_ind1, $inputs, 1, $dϕmvecs, true, $hw_channel) -# #@profview_allocs loop_forwarddiff_grad_nelboi_z(10_000, h21, approx, rnormM1, nljoint_args_ind1, 1, ϕm_, ϕqIc_, θsP_, dϕmvecs, dθsP, hw_channel, nothing, cor_ends.M, transM, nljoint_args_fix, n_MC_val) - -# approx2 = MeanHVIApproximation() -# hiM1 = h0M.helpers_sites[1] -# @test isfinite(CP.compute_nelboi_z!(hiM1, 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, -# ϕms1, ϕqIc, θsP1, cor_endsM, transM, nljoint_args_fix)) == 0 -# end -# alloc_compute_nelboi_zM!(hiM1, approx2, rnormM1, nljoint_args_ind1, -# ϕms1, ϕqIc, θsP1, cor_ends.M, transM, nljoint_args_fix) -# end - @testset "grad_neg_elbo_sites" begin ϕqIc = intϕqI(ϕqI) ϕqPc = intϕqP(ϕqP) @@ -237,13 +135,30 @@ pbm_covar_indices = SA.SVector{length(pbm_covars)}([findfirst(==(k), keys(par_te @test (@allocated loop_grad_neg_elbo_sites(100,h0, gradh0, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix; n_MC_val)) <= 6_748_800 end alloc_grad_neg_elbo_sites(h0, gradh0, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix; n_MC_val) - #_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) end @testset "HybridPointSolverSite" begin - + hp = CP.prepare_pointloss_helpers(approx, ϕg, ϕqP; n_θP, n_θM, n_site, n_MC, n_cov, n_covP, n_M, cor_ends, n_obs); + indiv_args = indiv_args1 + CP.check_pointloss_helpers(hp, indiv_args.xM, pbm_covar_indices) + ϕ = (;ϕg=ϕgv, ϕqP, ϕqI) + sample_args = (; + approx, g, is_testmode=true, pbm_covar_indices, intϕqP, intϕqI, cor_ends, + transP, transM) + nL = CP.neg_pointloss_sites!(hp, ϕ, sample_args, indiv_args, nljoint_args_fix) + @test isfinite(nL) + function loop_neg_pointloss_sites!(n, hp, ϕ, sample_args, indiv_args, nljoint_args_fix) + for _ in 1:n + @noinline CP.neg_pointloss_sites!(hp, ϕ, sample_args, indiv_args, nljoint_args_fix) + end + nothing + end + # allocations in StatefulLuxLayer + #@profview_allocs loop_neg_pointloss_sites!(100_000, hp, ϕ, sample_args, indiv_args, nljoint_args_fix) + function alloc_neg_pointloss_sites!(hp, ϕ, sample_args, indiv_args, nljoint_args_fix) + @allocated loop_neg_pointloss_sites!(10, hp, ϕ, sample_args, indiv_args, nljoint_args_fix) + end + alloc_neg_pointloss_sites!(hp, ϕ, sample_args, indiv_args, nljoint_args_fix) <= 47_200 end From 2ef9041c9f3806d1124f0d09fe0197cde8c7bfca Mon Sep 17 00:00:00 2001 From: Thomas Wutzler Date: Wed, 21 Oct 2026 08:49:25 +0000 Subject: [PATCH 3/4] work on pointloss --- src/HybridVariationalInference.jl | 1 + src/pointloss_site.jl | 7 +- src/pointloss_site_grad.jl | 123 +++++++++++++++--------------- test/test_elbo_site_doubleMM.jl | 20 +++-- 4 files changed, 78 insertions(+), 73 deletions(-) diff --git a/src/HybridVariationalInference.jl b/src/HybridVariationalInference.jl index 9f1d883..905a0d7 100644 --- a/src/HybridVariationalInference.jl +++ b/src/HybridVariationalInference.jl @@ -170,6 +170,7 @@ include("elbo_site_corr.jl") #export prepare_elbo_helpers include("pointloss_site.jl") +include("pointloss_site_grad.jl") #export get_loss_ran_tr_f diff --git a/src/pointloss_site.jl b/src/pointloss_site.jl index 259ad8c..1d98b16 100644 --- a/src/pointloss_site.jl +++ b/src/pointloss_site.jl @@ -3,7 +3,7 @@ elbo_helpers need to be initialized with new random numbers in h.ζP and hi.ζM by calling randnPM! before. By this way, we can compute the derivative corresponding to the forward pass """ -function neg_pointloss_sites!( +function pointloss_sites!( pointloss_helpers::NamedTuple, # tuple of preallocated arrays ϕ::NamedTuple, sample_args::NamedTuple, @@ -25,9 +25,8 @@ function neg_pointloss_sites!( # h.ζP .= CA.getdata(view(ϕqPc,Val(:μζP))) # (n_M x n_site) - # TODO g_apply without n_MC - ϕms = isnothing(pbm_covar_indices) ? view(h.ϕms,:,:,1) : h.ϕms - g_apply!(ϕms, ϕg, xM, h.ζP, pbm_covar_indices, g, h.xMP, is_testmode) + h_ϕms = isnothing(pbm_covar_indices) ? view(h.ϕms,:,:,1) : h.ϕms + g_apply!(h_ϕms, ϕg, xM, h.ζP, pbm_covar_indices, g, h.xMP, is_testmode) exp_ladJacTP = transformζ!(h.θP, transP, h.ζP) # return value captures ladJacT # so that one can provide its gradient to the pullback ϕm_it = eachcol(view(h.ϕms,:,:,1)) diff --git a/src/pointloss_site_grad.jl b/src/pointloss_site_grad.jl index 29c80d2..a7b6260 100644 --- a/src/pointloss_site_grad.jl +++ b/src/pointloss_site_grad.jl @@ -1,13 +1,11 @@ -function grad_neg_elbo_sites( - elbo_helpers::NamedTuple, # tuple of preallocated arrays - grad_elbo_helpers::NamedTuple, # tuple of preallocated arrays pullback closures +function grad_pointloss_sites( + pointloss_helpers::NamedTuple, # tuple of preallocated arrays + grad_pointloss_helpers::NamedTuple, # tuple of preallocated arrays pullback closures ϕ::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 # approx::AbstractHVIApproximation, # rnormPM::NamedTuple, # tuple of random numbers # ϕg::AbstractVector{TG}, ϕqP::AbstractVector{TF}, ϕqI::AbstractVector{TF}, g, @@ -21,8 +19,8 @@ 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 - h = elbo_helpers # preallocated μζP, dμζP, ζsP, ϕms, xMP, dxMP +) + h = pointloss_helpers # preallocated μζP, dμζP, ζsP, ϕms, xMP, dxMP (;ϕg, ϕqP, ϕqI) = ϕ TF = eltype(ϕqP) (;approx, g, is_testmode, pbm_covar_indices, intϕqP, intϕqI, cor_ends, transP, transM) = sample_args @@ -32,40 +30,44 @@ function grad_neg_elbo_sites( use_ϕm_matrix = isnothing(pbm_covar_indices) ϕqPc = intϕqP(ϕqP) ϕqIc = intϕqI(ϕqI) - ϕm_buffer_key = use_ϕm_matrix ? :ϕms : :ϕms_mcs - h_ϕm = h[ϕm_buffer_key] - check_elbo_helpers(h, xM, pbm_covar_indices; n_ϕg = length(ϕg)) + #ϕm_buffer_key = use_ϕm_matrix ? :ϕms : :ϕms_mcs + #h_ϕms = h[ϕm_buffer_key] + h_ϕms = isnothing(pbm_covar_indices) ? view(h.ϕms,:,:,1) : h.ϕms + check_pointloss_helpers(h, xM, pbm_covar_indices) n_cov, n_site = size(xM) - n_θP, n_MC = size(h.ζsP) - @assert n_MC == NMC + n_θP, n_MC = size(h.ζP) + @assert n_MC == 1 ax_θP = get_axis_θP(nljoint_args_fix.f!) ax_θM = get_axis_θM(nljoint_args_fix.f!) + n_θM = axis_length(ax_θM) nljoint_args_indi = NamedTuple{keys(nljoint_args_inds)}( first(zip_eachlastdims(nljoint_args_inds))) - gradh = !isempty(grad_elbo_helpers) ? grad_elbo_helpers : + gradh = !isempty(grad_pointloss_helpers) ? grad_pointloss_helpers : #Test.@inferred prepare_gradelbo_helpers( prepare_gradelbo_helpers( - ϕqIc, selectdim(h_ϕm, ndims(h_ϕm), 1), h.ζsP, ϕg, ϕqPc, approx; + ϕqIc, selectdim(h_ϕms, ndims(h_ϕms), 1), h.ζP, ϕ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) + h, rnormMi = Zeros(TF, n_θM, n_MC), nljoint_args_indi, + n_site, n_cov, cor_ends, transP, transM, nljoint_args_fix, n_MC_val=Val(1)) hw_channel = gradh.hw_channel + Main.@infiltrate_main # - sample_ζsP!(h.ζsP, h.logσ_ζP, 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 + h.ζP .= CA.getdata(view(ϕqPc,Val(:μζP))) + g_apply!(h_ϕms, ϕg, xM, h.ζP, pbm_covar_indices, g, h.xMP, is_testmode) + ladJacTP = transformζ!(h.θsP, transP, h.ζP) # return value captures ladJacT # # parallel ForwardDiffGradient through forwarddiff_grad_nelboi_z! - cl = ForwardDiffGradNelboiZCl( + cl = ForwardDiffGradPointlossZCl( approx, ϕqIc, h.θsP, gradh.dϕmvecs, gradh.dθsP, gradh.hw_channel, cor_ends.M, transM, Val(keys(nljoint_args_inds)), nljoint_args_fix, n_MC_val) - ϕm_it = eachslice(h_ϕm; dims = ndims(h_ϕm)) + ϕm_it = eachcol(view(h.ϕms,:,:,1)) init = (; dϕqI = zero(TF) .* SA.SVector{axis_length(CA.getaxes(ϕqIc)[1])}(ϕqIc), dθsP = zero(TF) .* SA.SVector{axis_length(ax_θP)*NMC}(h.θsP), ) res = Folds.mapreduce(cl, make_tuple_reducer(+), - zip(h.helpers_sites, rnormPM.M, zip_eachlastdims(nljoint_args_inds), ϕm_it, 1:n_site), + zip(h.helpers_sites, zip_eachlastdims(nljoint_args_inds), ϕm_it, 1:n_site), executor; init) + error("implement pullbacks") # copy SharedArray into preallocated plain array for ∂elbo_∂θP ∂elbo_∂ϕqm = gradh.∂elbo_∂ϕqm copyto!(∂elbo_∂ϕqm, gradh.dϕmvecs) # copy from shared array @@ -79,8 +81,8 @@ function grad_neg_elbo_sites( # # pullback gradients of ϕqm -> gradh.dϕg and gradh.dζsP gradh.pullback_g_apply!( - gradh.dϕg, gradh.∂elbo_∂ϕm_∂ζP, h_ϕm, ∂elbo_∂ϕqm, - ϕg, xM, h.ζsP, pbm_covar_indices, g, is_testmode) + gradh.dϕg, gradh.∂elbo_∂ϕm_∂ζP, h_ϕms, ∂elbo_∂ϕqm, + ϕg, xM, h.ζP, pbm_covar_indices, g, is_testmode) # # pullback gradients of ∂elbo_∂θP to ∂elbo_∂θP_∂ζP ∂elbo_∂θP_∂ζP = gradh.∂elbo_∂θP_∂ζP @@ -90,14 +92,14 @@ function grad_neg_elbo_sites( ∂elbo_∂ladJacTP, h.θsP, transP, - h.ζsP, + h.ζP, ) # # pullback gradients of ∂elbo_∂ζP, ∂elbo_∂ϕm_∂ζP, and ∂elbo_∂logσ_ζP to dϕqP gradh.pullback_cl_sample_ζsP!( gradh.dϕqPc, ∂elbo_∂θP_∂ζP + gradh.∂elbo_∂ϕm_∂ζP, gradh.∂elbo_∂logσ_ζP, - h.ζsP, h.logσ_ζP, approx, rnormPM.P, ϕqPc, cor_ends.P, h.sample_buffers, + h.ζP, 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 @@ -105,7 +107,7 @@ end """ Callable to make deliver arguments that do not differ by individual to Foldl.mapreduce. """ -struct ForwardDiffGradNelboiZCl{KEYS, TA, Tϕq, Tθ, TD, TP, THWC, TC, TM, TAF, NMC} +struct ForwardDiffGradPointlossiZCl{KEYS, TA, Tϕq, Tθ, TD, TP, THWC, TC, TM, TAF, NMC} approx::TA ϕqIc::Tϕq θsP::Tθ @@ -119,13 +121,14 @@ struct ForwardDiffGradNelboiZCl{KEYS, TA, Tϕq, Tθ, TD, TP, THWC, TC, TM, TAF, nljoint_args_fix::TAF n_MC_val::Val{NMC} end -function (f::ForwardDiffGradNelboiZCl{KEYS})(tup) where KEYS - hi, rnormM, nljoint_args_indi_tup, ϕm, i = tup +function (f::ForwardDiffGradPointlossiZCl{KEYS})(tup) where KEYS + local hi, nljoint_args_indi_tup, ϕm, i + hi, nljoint_args_indi_tup, ϕm, i = tup #nljoint_args_indi = NamedTuple{f.keys_nljoint_args_indi}(nljoint_args_indi_tup) nljoint_args_indi = NamedTuple{KEYS}(nljoint_args_indi_tup) - forwarddiff_grad_nelboi_z!(hi, f.approx, rnormM, nljoint_args_indi, ϕm, i, + forwarddiff_grad_pointloss_z!(hi, f.approx, nljoint_args_indi, ϕm, i, f.ϕqIc, f.θsP, f.dϕmvecs, f.dθsP, f.hw_channel, f.corendsM, f.transM, - f.nljoint_args_fix, f.n_MC_val, + f.nljoint_args_fix, nothing, #true, ) @@ -134,11 +137,11 @@ end # ϕm vector version # TOOD pass closure object rather than many parameters -function forwarddiff_grad_nelboi_z!(hi, approx::AbstractHVIApproximation, rnormM, +function forwarddiff_grad_pointloss_z!(hi, approx::AbstractHVIApproximation, nljoint_args_indi, ϕm::AbstractVector{TF}, i, ϕqIc, θsP, dϕmvecs, dθsP, hw_channel::Channel, cor_endsM, transM::Stacked, - nljoint_args_fix, n_MC_val::Val{NMC}, - omit_gradient=nothing) where {TF, NMC} + nljoint_args_fix, + omit_gradient=nothing) where {TF} # need to calls to ForwardDiff.gradient! to avoid reshape # group arrays of same dimensionality (here grad_ϕm and grad_ϕqI) into a Vcat ax_ϕqI = CA.getaxes(ϕqIc)[1] @@ -146,7 +149,7 @@ function forwarddiff_grad_nelboi_z!(hi, approx::AbstractHVIApproximation, rnormM with_channel_element(hw_channel) do hwi_tup hwi, hwi_dc = hwi_tup #grad_conf = hwi.grad_conf - nelboi_z_mat, nelboi_z_vec = make_nelboiz_lazy_cls(hi, hwi_dc, approx, rnormM, nljoint_args_indi, + nelboi_z_mat, nelboi_z_vec = make_pointlossz_lazy_cls(hi, hwi_dc, approx, nljoint_args_indi, ϕm, ϕqIc, θsP, cor_endsM, transM, nljoint_args_fix, ax_ϕqI, Val(1), @@ -170,22 +173,18 @@ function forwarddiff_grad_nelboi_z!(hi, approx::AbstractHVIApproximation, rnormM local res res = (; dϕqI = SA.SVector{axis_length(CA.getaxes(ϕqIc)[1])}(grad_ϕqI), - dθsP = SA.SVector{axis_length(ax_ϕP)*NMC}(grad_θsP), + dθsP = SA.SVector{axis_length(ax_ϕP)*1}(grad_θsP), ) end end -#@noinline CP.forwarddiff_grad_nelboi_z!(hi, approx, rnormM, nljoint_args_indi, ϕm_, i, -# ϕqIc_, θsP_, dϕmvecs_, dθsP_, hw_channel, cor_endsM, transM, nljoint_args_fix, omit_gradient) - - # ϕm matrix version -function forwarddiff_grad_nelboi_z!(hi, approx::AbstractHVIApproximation, rnormM, +function forwarddiff_grad_pointloss_z!(hi, approx::AbstractHVIApproximation, nljoint_args_indi, ϕm::AbstractMatrix{TF}, i, ϕqIc, θsP, dϕmvecs, dθsP, hw_channel::Channel, cor_endsM, transM::Stacked, - nljoint_args_fix, n_MC_val::Val{NMC}, - omit_gradient=nothing) where {TF, NMC} + nljoint_args_fix, + omit_gradient=nothing) where {TF} # need to calls to ForwardDiff.gradient! to avoid reshape # group arrays of same dimensionality (here grad_ϕm and grad_ϕqI) into a Vcat ax_ϕqI = CA.getaxes(ϕqIc)[1] @@ -193,7 +192,7 @@ function forwarddiff_grad_nelboi_z!(hi, approx::AbstractHVIApproximation, rnormM with_channel_element(hw_channel) do hwi_tup hwi, hwi_dc = hwi_tup #grad_conf = hwi.grad_conf - nelboi_z_mat, nelboi_z_vec = make_nelboiz_lazy_cls(hi, hwi_dc, approx, rnormM, nljoint_args_indi, + nelboi_z_mat, nelboi_z_vec = make_pointlossz_lazy_cls(hi, hwi_dc, approx, nljoint_args_indi, ϕm, ϕqIc, θsP, cor_endsM, transM, nljoint_args_fix, ax_ϕqI, Val(2), ## 2-dimensional ϕm @@ -222,13 +221,13 @@ function forwarddiff_grad_nelboi_z!(hi, approx::AbstractHVIApproximation, rnormM local res res = (; dϕqI = SA.SVector{axis_length(CA.getaxes(ϕqIc)[1])}(grad_ϕqI), - dθsP = SA.SVector{axis_length(ax_ϕqP)*NMC}(grad_θsP), + dθsP = SA.SVector{axis_length(ax_ϕqP)*1}(grad_θsP), ) end end -function make_nelboiz_lazy_cls(hi, hwi_dc, approx::AbstractHVIApproximation, rnormM, nljoint_args_indi, +function make_pointlossz_lazy_cls(hi, hwi_dc, approx::AbstractHVIApproximation, nljoint_args_indi, ϕm, ϕqIc, θsP, cor_endsM, transM, nljoint_args_fix, ax_ϕqI::CA.Axis, ndim_ϕ::Val{1}, @@ -242,8 +241,8 @@ function make_nelboiz_lazy_cls(hi, hwi_dc, approx::AbstractHVIApproximation, rno ϕm_dual = copyto!(PAT.get_tmp(hwi_dc.ϕm_dc, mat), ϕm) ϕqIc_dual = CA.ComponentArray( copyto!(PAT.get_tmp(hwi_dc.ϕqI_dc, mat), ϕqIc), CA.getaxes(ϕqIc)) - compute_nelboi_z!( - hi, approx, rnormM, nljoint_args_indi, + compute_pointlossi_z!( + hi, nljoint_args_indi, ϕm_dual, ϕqIc_dual, θsP_dual, cor_endsM, transM, nljoint_args_fix, )[1] @@ -253,8 +252,8 @@ function make_nelboiz_lazy_cls(hi, hwi_dc, approx::AbstractHVIApproximation, rno ϕm_dual = vec.args[1] ϕqIc_dual = CA.ComponentVector(vec.args[2], ax_ϕqI) θsP_dual = copyto!(PAT.get_tmp(hwi_dc.θsP_dc, vec), θsP) # reshape(reinterpred( - compute_nelboi_z!( - hi, approx, rnormM, nljoint_args_indi, + compute_pointlossi_z!( + hi, nljoint_args_indi, ϕm_dual, ϕqIc_dual, θsP_dual, cor_endsM, transM, nljoint_args_fix, )[1] @@ -263,7 +262,7 @@ function make_nelboiz_lazy_cls(hi, hwi_dc, approx::AbstractHVIApproximation, rno end # ϕm Matrix case -function make_nelboiz_lazy_cls(hi, hwi_dc, approx::AbstractHVIApproximation, rnormM, nljoint_args_indi, +function make_pointlossz_lazy_cls(hi, hwi_dc, approx::AbstractHVIApproximation, rnormM, nljoint_args_indi, ϕm, ϕqIc, θsP, cor_endsM, transM, nljoint_args_fix, ax_ϕqI::CA.Axis, ndim_ϕ::Val{2}, @@ -276,8 +275,8 @@ function make_nelboiz_lazy_cls(hi, hwi_dc, approx::AbstractHVIApproximation, rno ϕm_dual = mat.args[2] ϕqIc_dual = CA.ComponentArray( copyto!(PAT.get_tmp(hwi_dc.ϕqI_dc, mat), ϕqIc), CA.getaxes(ϕqIc)) - compute_nelboi_z!( - hi, approx, rnormM, nljoint_args_indi, + compute_pointlossi_z!( + hi, nljoint_args_indi, ϕm_dual, ϕqIc_dual, θsP_dual, cor_endsM, transM, nljoint_args_fix, )[1] @@ -287,8 +286,8 @@ function make_nelboiz_lazy_cls(hi, hwi_dc, approx::AbstractHVIApproximation, rno ϕqIc_dual = CA.ComponentVector(vec, ax_ϕqI) ϕm_dual = copyto!(PAT.get_tmp(hwi_dc.ϕm_dc, vec), ϕm) θsP_dual = copyto!(PAT.get_tmp(hwi_dc.θsP_dc, vec), θsP) - compute_nelboi_z!( - hi, approx, rnormM, nljoint_args_indi, + compute_pointlossi_z!( + hi, nljoint_args_indi, ϕm_dual, ϕqIc_dual, θsP_dual, cor_endsM, transM, nljoint_args_fix, )[1] @@ -480,8 +479,8 @@ function prepare_gradelbo_helpers( # ϕqIc = cv_grad[Val(:ϕqIc)] # ϕm = cv_grad[Val(:ϕm)] # θsP = cv_grad[Val(:θsP)] - use_ϕm_vector = isnothing(pbm_covar_indices) # for one site, provide singe ϕm across MC - n_covP = use_ϕm_vector ? 0 : length(pbm_covar_indices) + use_ϕm_vector = ND==1 # single ϕm across MC (for each site) + n_covP = isnothing(pbm_covar_indices) ? 0 : length(pbm_covar_indices) n_ϕmvec = length(ϕm) #use_ϕm_matrix ? n_M : n_M * n_MC n_θP, n_MC = size(θsP) @assert n_MC == NMC @@ -501,7 +500,7 @@ function prepare_gradelbo_helpers( ) # make sure, each worker gets its own copy, use similar hw_dcs = (hw_dcs_raw) -> map_leaves_nt(x -> PAT.DiffCache(similar(x), chunk), hw_dcs_raw) - nelboi_z_mat, nelboi_z_vec = make_nelboiz_lazy_cls(hi1, hw_dcs(hw_dcs_raw), approx, rnormMi, nljoint_args_indi, + nelboi_z_mat, nelboi_z_vec = make_pointlossz_lazy_cls(hi1, hw_dcs(hw_dcs_raw), approx, rnormMi, nljoint_args_indi, ϕm, ϕqIc, θsP, cor_ends.M, transM, nljoint_args_fix, CA.getaxes(ϕqIc)[1], Val(ND), @@ -520,16 +519,16 @@ function prepare_gradelbo_helpers( end grad_conf_mat_ϕmvec, grad_conf_vec_ϕmvec else - local rad_θP_ϕm1, diffchunks_vec, diffchunks_mat + local grad_θP_ϕm1, diffchunks_vec, diffchunks_mat diffchunks_vec = ForwardDiff.Chunk(min(chunk, axis_length(CA.getaxes(ϕqIc)[1]))) diffchunks_mat = ForwardDiff.Chunk(min(chunk, NMC*(axis_length(CA.getaxes(ϕqPc)[1]) + axis_length(ax_θM)))) grad_conf_vec_ϕmmat() = ForwardDiff.GradientConfig(nelboi_z_vec, CA.getdata(ϕqIc), diffchunks_vec) - rad_θP_ϕm1 = grad_θsP_ϕm(ϕm, θsP) + grad_θP_ϕm1 = grad_θsP_ϕm(ϕm, θsP) function grad_conf_mat_ϕmmat() - cfg_mat = ForwardDiff.GradientConfig(nelboi_z_mat, rad_θP_ϕm1, diffchunks_mat) - update_duals_Vcat(cfg_mat, rad_θP_ϕm1) + cfg_mat = ForwardDiff.GradientConfig(nelboi_z_mat, grad_θP_ϕm1, diffchunks_mat) + update_duals_Vcat(cfg_mat, grad_θP_ϕm1) end grad_conf_mat_ϕmmat, grad_conf_vec_ϕmmat end diff --git a/test/test_elbo_site_doubleMM.jl b/test/test_elbo_site_doubleMM.jl index 34a62c4..7e36216 100644 --- a/test/test_elbo_site_doubleMM.jl +++ b/test/test_elbo_site_doubleMM.jl @@ -131,6 +131,7 @@ pbm_covar_indices = SA.SVector{length(pbm_covars)}([findfirst(==(k), keys(par_te function alloc_grad_neg_elbo_sites(h0, gradh0, ϕ, 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) + GC.gc() #@profview_allocs loop_grad_neg_elbo_sites(1_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 end @@ -145,20 +146,25 @@ end sample_args = (; approx, g, is_testmode=true, pbm_covar_indices, intϕqP, intϕqI, cor_ends, transP, transM) - nL = CP.neg_pointloss_sites!(hp, ϕ, sample_args, indiv_args, nljoint_args_fix) + nL = CP.pointloss_sites!(hp, ϕ, sample_args, indiv_args, nljoint_args_fix) @test isfinite(nL) - function loop_neg_pointloss_sites!(n, hp, ϕ, sample_args, indiv_args, nljoint_args_fix) + function loop_pointloss_sites!(n, hp, ϕ, sample_args, indiv_args, nljoint_args_fix) for _ in 1:n - @noinline CP.neg_pointloss_sites!(hp, ϕ, sample_args, indiv_args, nljoint_args_fix) + @noinline CP.pointloss_sites!(hp, ϕ, sample_args, indiv_args, nljoint_args_fix) end nothing end # allocations in StatefulLuxLayer - #@profview_allocs loop_neg_pointloss_sites!(100_000, hp, ϕ, sample_args, indiv_args, nljoint_args_fix) - function alloc_neg_pointloss_sites!(hp, ϕ, sample_args, indiv_args, nljoint_args_fix) - @allocated loop_neg_pointloss_sites!(10, hp, ϕ, sample_args, indiv_args, nljoint_args_fix) + #@profview_allocs loop_pointloss_sites!(100_000, hp, ϕ, sample_args, indiv_args, nljoint_args_fix) + function alloc_pointloss_sites!(hp, ϕ, sample_args, indiv_args, nljoint_args_fix) + @allocated loop_pointloss_sites!(10, hp, ϕ, sample_args, indiv_args, nljoint_args_fix) end - alloc_neg_pointloss_sites!(hp, ϕ, sample_args, indiv_args, nljoint_args_fix) <= 47_200 + alloc_pointloss_sites!(hp, ϕ, sample_args, indiv_args, nljoint_args_fix) <= 47_200 + # + # res0, gradh0 = CP.grad_pointloss_sites( + # #@descend_code_warntype CP.neg_elbo_sites!( + # hp, (;), ϕ, sample_args, indiv_args, nljoint_args_fix) + end From ba130a28373627e37eaf5c249dcd77978ebd818c Mon Sep 17 00:00:00 2001 From: Thomas Wutzler Date: Wed, 21 Oct 2026 15:44:53 +0000 Subject: [PATCH 4/4] implement grad_pointloss_sites --- src/elbo_site.jl | 4 +- src/elbo_site_grad.jl | 19 +- src/pointloss_site.jl | 62 ++-- src/pointloss_site_grad.jl | 487 ++++++-------------------------- test/runtests.jl | 1 + test/test_elbo_site.jl | 2 +- test/test_elbo_site_doubleMM.jl | 34 ++- 7 files changed, 172 insertions(+), 437 deletions(-) diff --git a/src/elbo_site.jl b/src/elbo_site.jl index 3a77312..04bc3f3 100644 --- a/src/elbo_site.jl +++ b/src/elbo_site.jl @@ -28,8 +28,8 @@ function neg_elbo_sites!( nljoint_args_inds = NamedTuple{keys(indiv_args)[r2end]}(values(indiv_args)[r2end]) @assert size(rnormPM.P) == size(h.ζsP) use_dc = h.helpers_sites[1].ζsM isa PAT.DiffCache - n_M, n_MC = use_dc ? size(h.helpers_sites[1].ζsM.du) : size(h.helpers_sites[1].ζsM) - @assert size(rnormPM.M[1]) == (n_M, n_MC) + n_θM, n_MC = use_dc ? size(h.helpers_sites[1].ζsM.du) : size(h.helpers_sites[1].ζsM) + @assert size(rnormPM.M[1]) == (n_θM, n_MC) ϕqPc = intϕqP(ϕqP) ϕqIc = intϕqI(ϕqI) pbm_covar_indices = !isnothing(pbm_covar_indices) && isempty(pbm_covar_indices) ? nothing : pbm_covar_indices diff --git a/src/elbo_site_grad.jl b/src/elbo_site_grad.jl index 29c80d2..919eb7e 100644 --- a/src/elbo_site_grad.jl +++ b/src/elbo_site_grad.jl @@ -516,7 +516,7 @@ function prepare_gradelbo_helpers( grad_ϕm_ϕqI1 = grad_ϕm_ϕqI(ϕm, ϕqIc) function grad_conf_vec_ϕmvec() cfg_vec = ForwardDiff.GradientConfig(nelboi_z_vec, grad_ϕm_ϕqI1, diffchunk_vec) - update_duals_Vcat(cfg_vec, grad_ϕm_ϕqI1) + update_duals_Vcat2(cfg_vec, grad_ϕm_ϕqI1) end grad_conf_mat_ϕmvec, grad_conf_vec_ϕmvec else @@ -529,7 +529,7 @@ function prepare_gradelbo_helpers( rad_θP_ϕm1 = grad_θsP_ϕm(ϕm, θsP) function grad_conf_mat_ϕmmat() cfg_mat = ForwardDiff.GradientConfig(nelboi_z_mat, rad_θP_ϕm1, diffchunks_mat) - update_duals_Vcat(cfg_mat, rad_θP_ϕm1) + update_duals_Vcat2(cfg_mat, rad_θP_ϕm1) end grad_conf_mat_ϕmmat, grad_conf_vec_ϕmmat end @@ -579,15 +579,28 @@ end @inline grad_θsP_ϕm(ϕm::AbstractMatrix, θsP) = Vcat(similar(θsP), similar(ϕm)) -@inline function update_duals_Vcat(cfg::ForwardDiff.GradientConfig{T,V,N}, vc::Vcat) where {T,V,N} +@inline function update_duals_Vcat2(cfg::ForwardDiff.GradientConfig{T,V,N}, vc::Vcat) where {T,V,N} # similar Vcat returns a vector, need to convert dual of cfg to Vcat # no view, but copy with index_at_dim + @assert length(vc.args) == 2 duals_vcat = Vcat(index_at_dim(cfg.duals, axes(vc.args[1], 1); dim=1), index_at_dim(cfg.duals, axes(vc.args[2], 1); dim=1)); cfg = ForwardDiff.GradientConfig{T,V,N,typeof(duals_vcat)}(cfg.seeds, duals_vcat) @assert cfg.duals isa Vcat cfg end +@inline function update_duals_Vcat3(cfg::ForwardDiff.GradientConfig{T,V,N}, vc::Vcat) where {T,V,N} + # similar Vcat returns a vector, need to convert dual of cfg to Vcat + # no view, but copy with index_at_dim + @assert length(vc.args) == 3 + duals_vcat = Vcat(index_at_dim(cfg.duals, axes(vc.args[1], 1); dim=1), + index_at_dim(cfg.duals, axes(vc.args[2], 1); dim=1), + index_at_dim(cfg.duals, axes(vc.args[3], 1); dim=1), + ); + cfg = ForwardDiff.GradientConfig{T,V,N,typeof(duals_vcat)}(cfg.seeds, duals_vcat) + @assert cfg.duals isa Vcat + cfg +end diff --git a/src/pointloss_site.jl b/src/pointloss_site.jl index 1d98b16..4831fd4 100644 --- a/src/pointloss_site.jl +++ b/src/pointloss_site.jl @@ -17,28 +17,28 @@ function pointloss_sites!( @assert keys(indiv_args)[1] == :xM xM = indiv_args.xM; r2end = 2:length(indiv_args) nljoint_args_inds = NamedTuple{keys(indiv_args)[r2end]}(values(indiv_args)[r2end]) - use_dc = h.helpers_sites[1].ζM isa PAT.DiffCache - (n_M,) = use_dc ? size(h.helpers_sites[1].ζM.du) : size(h.helpers_sites[1].ζM) + use_dc = h.helpers_sites[1].ζMm isa PAT.DiffCache + (n_θM,) = use_dc ? size(h.helpers_sites[1].ζMm.du) : size(h.helpers_sites[1].ζMm) ϕqPc = intϕqP(ϕqP) ϕqIc = intϕqI(ϕqI) pbm_covar_indices = !isnothing(pbm_covar_indices) && isempty(pbm_covar_indices) ? nothing : pbm_covar_indices # - h.ζP .= CA.getdata(view(ϕqPc,Val(:μζP))) + h.ζPm .= CA.getdata(view(ϕqPc,Val(:μζP))) # (n_M x n_site) - h_ϕms = isnothing(pbm_covar_indices) ? view(h.ϕms,:,:,1) : h.ϕms - g_apply!(h_ϕms, ϕg, xM, h.ζP, pbm_covar_indices, g, h.xMP, is_testmode) - exp_ladJacTP = transformζ!(h.θP, transP, h.ζP) # return value captures ladJacT + h_ϕms = isnothing(pbm_covar_indices) ? view(h.ϕmsa,:,1,:) : h.ϕmsa + g_apply!(h_ϕms, ϕg, xM, h.ζPm, pbm_covar_indices, g, h.xMP, is_testmode) + exp_ladJacTP = transformζ!(h.θPm, transP, h.ζPm) # return value captures ladJacT # so that one can provide its gradient to the pullback - ϕm_it = eachcol(view(h.ϕms,:,:,1)) + ϕm_it = eachslice(h.ϕmsa; dims=3) template = ϕqI # only important for gradient - θP = h.θP + θPm = h.θPm keys_nljoint_args_indi_val = Val(keys(nljoint_args_inds)) # closure with approx, transM, and kwargs function compute_pointlossi_z_cl!(hi, 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_pointlossi_z!(hi, nljoint_args_indi, ϕm, - ϕqIc, θP, cor_ends.M, transM, nljoint_args_fix) + ϕqIc, θPm, cor_ends.M, transM, nljoint_args_fix) end #res_site = map(compute_pointlossi_z!, h.helpers_sites, rnormPM.M, i_sites_train, ϕm_it) #MAYBE: distributed mapreduce: @@ -51,20 +51,20 @@ function pointloss_sites!( end function compute_pointlossi_z!(hi, - nljoint_args_indi, ϕm::AbstractArray, ϕqIc::AbstractArray, θP, + nljoint_args_indi, ϕms::AbstractMatrix, ϕqIc::AbstractVector, θPm::AbstractMatrix, cor_endsM, transM, nljoint_args_fix, ) # DM is 1 (vector) if no population covariates - constant across MC # 2 (Matrix) if g is driven by population covariates - differs across MC # on update -> sync corresponding function within grad_neg_elbo_sites - if hi.ζM isa PAT.DiffCache + if hi.ζMm isa PAT.DiffCache # Gradiend called twice because of different dimensions need to pick the Dual - template = eltype(ϕqIc) <: ForwardDiff.Dual ? ϕqIc : θP + template = eltype(ϕqIc) <: ForwardDiff.Dual ? ϕqIc : θPm hi = map_leaves_nt(x -> PAT.get_tmp(x, template), hi) end n_θM = cor_endsM[end] - hi.ζM .= view(ϕm, 1:n_θM, :) - transformζ!(hi.θM, transM, hi.ζM) - exp_nL = exp_nLi!(hi.y, θP, hi.θM, nljoint_args_indi, nljoint_args_fix)[1] + hi.ζMm .= view(ϕms, 1:n_θM, :) # μζM in first n_θM entries + transformζ!(hi.θMm, transM, hi.ζMm) + exp_nL = exp_nLi!(hi.ym, θPm, hi.θMm, nljoint_args_indi, nljoint_args_fix)[1] pointlossi = exp_nL end @@ -75,15 +75,15 @@ function prepare_pointloss_helpers(approx::AbstractHVIApproximation, diffchunk::ForwardDiff.Chunk{chunk} = ForwardDiff.Chunk(8), ) where {TG, TF, use_dc, chunk} his = Tuple((; - ζM = Matrix{TF}(undef, n_θM, 1), - θM = Matrix{TF}(undef, n_θM, 1), - y = Matrix{TF}(undef, n_obs,1), + ζMm = Matrix{TF}(undef, n_θM, 1), + θMm = Matrix{TF}(undef, n_θM, 1), + ym = Matrix{TF}(undef, n_obs,1), ) for i in 1:n_site) helpers_sites = use_dc ? map_leaves_nt(x -> PAT.DiffCache(x, chunk), his) : his h = (; - ζP = Matrix{TF}(undef, n_θP,1), - θP = Matrix{TF}(undef, n_θP,1), - ϕms = Array{TF}(undef, n_M, n_site,1), + ζPm = Matrix{TF}(undef, n_θP,1), + θPm = Matrix{TF}(undef, n_θP,1), + ϕmsa = Array{TF}(undef, n_M, 1, n_site), xMP = Matrix{TG}(undef, (n_cov + n_covP), n_site), helpers_sites, ) @@ -92,20 +92,20 @@ end function check_pointloss_helpers(h::NamedTuple, xM::AbstractMatrix, pbm_covar_indices) n_cov, n_site = size(xM) n_covP = isnothing(pbm_covar_indices) ? 0 : length(pbm_covar_indices) - (n_θP,) = size(h.ζP) - n_M = size(h.ϕms, 1) - @assert size(h.ζP) == (n_θP, 1) - @assert size(h.θP) == (n_θP, 1) - @assert size(h.ϕms) == (n_M, n_site,1) + (n_θP,_) = size(h.ζPm) + n_M = size(h.ϕmsa, 1) + @assert size(h.ζPm) == (n_θP, 1) + @assert size(h.θPm) == (n_θP, 1) + @assert size(h.ϕmsa) == (n_M, 1, n_site) @assert size(h.xMP) == ((n_cov + n_covP), n_site) # @assert length(h.helpers_sites) == n_site hi = h.helpers_sites[1] - n_θM = size(hi.ζM.du, 1) - n_obs = size(hi.y.du, 1) - @assert size(hi.ζM.du) == (n_θM, 1) - @assert size(hi.θM.du) == (n_θM, 1) - @assert size(hi.y.du) == (n_obs, 1) + n_θM = size(hi.ζMm.du, 1) + n_obs = size(hi.ym.du, 1) + @assert size(hi.ζMm.du) == (n_θM, 1) + @assert size(hi.θMm.du) == (n_θM, 1) + @assert size(hi.ym.du) == (n_obs, 1) end diff --git a/src/pointloss_site_grad.jl b/src/pointloss_site_grad.jl index a7b6260..5fc157a 100644 --- a/src/pointloss_site_grad.jl +++ b/src/pointloss_site_grad.jl @@ -32,42 +32,41 @@ function grad_pointloss_sites( ϕqIc = intϕqI(ϕqI) #ϕm_buffer_key = use_ϕm_matrix ? :ϕms : :ϕms_mcs #h_ϕms = h[ϕm_buffer_key] - h_ϕms = isnothing(pbm_covar_indices) ? view(h.ϕms,:,:,1) : h.ϕms + h_ϕms = isnothing(pbm_covar_indices) ? view(h.ϕmsa,:,1,:) : h.ϕmsa check_pointloss_helpers(h, xM, pbm_covar_indices) n_cov, n_site = size(xM) - n_θP, n_MC = size(h.ζP) + n_θP, n_MC = size(h.ζPm) @assert n_MC == 1 ax_θP = get_axis_θP(nljoint_args_fix.f!) ax_θM = get_axis_θM(nljoint_args_fix.f!) n_θM = axis_length(ax_θM) + n_M = get_numberof_MLinputs(approx, 1:n_θM) nljoint_args_indi = NamedTuple{keys(nljoint_args_inds)}( first(zip_eachlastdims(nljoint_args_inds))) gradh = !isempty(grad_pointloss_helpers) ? grad_pointloss_helpers : - #Test.@inferred prepare_gradelbo_helpers( - prepare_gradelbo_helpers( - ϕqIc, selectdim(h_ϕms, ndims(h_ϕms), 1), h.ζP, ϕg, ϕqPc, approx; + Test.@inferred prepare_gradpointloss_helpers( + #prepare_gradpointloss_helpers( + ϕqIc, selectdim(h_ϕms, ndims(h_ϕms), 1), h.ζPm, ϕg, ϕqPc, approx; pbm_covar_indices, n_workers, - h, rnormMi = Zeros(TF, n_θM, n_MC), nljoint_args_indi, - n_site, n_cov, cor_ends, transP, transM, nljoint_args_fix, n_MC_val=Val(1)) - hw_channel = gradh.hw_channel - Main.@infiltrate_main + h, nljoint_args_indi, + n_site, n_M, n_cov, cor_ends, transP, transM, nljoint_args_fix, ) # - h.ζP .= CA.getdata(view(ϕqPc,Val(:μζP))) - g_apply!(h_ϕms, ϕg, xM, h.ζP, pbm_covar_indices, g, h.xMP, is_testmode) - ladJacTP = transformζ!(h.θsP, transP, h.ζP) # return value captures ladJacT + h.ζPm .= CA.getdata(view(ϕqPc,Val(:μζP))) + g_apply!(h_ϕms, ϕg, xM, h.ζPm, pbm_covar_indices, g, h.xMP, is_testmode) + ladJacTP = transformζ!(h.θPm, transP, h.ζPm) # return value captures ladJacT # - # parallel ForwardDiffGradient through forwarddiff_grad_nelboi_z! - cl = ForwardDiffGradPointlossZCl( - approx, ϕqIc, h.θsP, gradh.dϕmvecs, gradh.dθsP, gradh.hw_channel, cor_ends.M, transM, - Val(keys(nljoint_args_inds)), nljoint_args_fix, n_MC_val) - ϕm_it = eachcol(view(h.ϕms,:,:,1)) + # parallel ForwardDiffGradient through forwarddiff_grad_pointlossi_z! + cl = ForwardDiffGradPointlossiZCl( + approx, ϕqIc, h.θPm, gradh.dϕmvecs, gradh.hw_channel, cor_ends.M, transM, + Val(keys(nljoint_args_inds)), nljoint_args_fix) + ϕm_it = eachslice(h.ϕmsa; dims = ndims(h.ϕmsa)) init = (; dϕqI = zero(TF) .* SA.SVector{axis_length(CA.getaxes(ϕqIc)[1])}(ϕqIc), - dθsP = zero(TF) .* SA.SVector{axis_length(ax_θP)*NMC}(h.θsP), ) + dθsP = zero(TF) .* SA.SVector{axis_length(ax_θP)*1}(h.θPm), ) res = Folds.mapreduce(cl, make_tuple_reducer(+), zip(h.helpers_sites, zip_eachlastdims(nljoint_args_inds), ϕm_it, 1:n_site), executor; init) - error("implement pullbacks") + # # copy SharedArray into preallocated plain array for ∂elbo_∂θP ∂elbo_∂ϕqm = gradh.∂elbo_∂ϕqm copyto!(∂elbo_∂ϕqm, gradh.dϕmvecs) # copy from shared array @@ -82,7 +81,7 @@ function grad_pointloss_sites( # pullback gradients of ϕqm -> gradh.dϕg and gradh.dζsP gradh.pullback_g_apply!( gradh.dϕg, gradh.∂elbo_∂ϕm_∂ζP, h_ϕms, ∂elbo_∂ϕqm, - ϕg, xM, h.ζP, pbm_covar_indices, g, is_testmode) + ϕg, xM, h.ζPm, pbm_covar_indices, g, is_testmode) # # pullback gradients of ∂elbo_∂θP to ∂elbo_∂θP_∂ζP ∂elbo_∂θP_∂ζP = gradh.∂elbo_∂θP_∂ζP @@ -90,36 +89,32 @@ function grad_pointloss_sites( ∂elbo_∂θP_∂ζP, ∂elbo_∂θP, ∂elbo_∂ladJacTP, - h.θsP, + h.θPm, transP, - h.ζP, + h.ζPm, ) # - # pullback gradients of ∂elbo_∂ζP, ∂elbo_∂ϕm_∂ζP, and ∂elbo_∂logσ_ζP to dϕqP - gradh.pullback_cl_sample_ζsP!( - gradh.dϕqPc, - ∂elbo_∂θP_∂ζP + gradh.∂elbo_∂ϕm_∂ζP, gradh.∂elbo_∂logσ_ζP, - h.ζP, h.logσ_ζP, approx, rnormPM.P, ϕqPc, cor_ends.P, h.sample_buffers, - ) + # pullback gradients of ∂elbo_∂ζP, ∂elbo_∂ϕm_∂ζP to dϕqP + # for μζP (first n_θM entries) its the sum of the woggles, others are zero + gradh.dϕqPc[1:n_θM] .= ∂elbo_∂θP_∂ζP + gradh.∂elbo_∂ϕm_∂ζP + gradh.dϕqPc[(n_θM+1):end] .= zero(TF) (;dϕqPc = gradh.dϕqPc, dϕqIc = gradh.dϕqIc, dϕg = gradh.dϕg), gradh end """ Callable to make deliver arguments that do not differ by individual to Foldl.mapreduce. """ -struct ForwardDiffGradPointlossiZCl{KEYS, TA, Tϕq, Tθ, TD, TP, THWC, TC, TM, TAF, NMC} +struct ForwardDiffGradPointlossiZCl{KEYS, TA, TI, Tθ, TD, THWC, TC, TM, TAF} approx::TA - ϕqIc::Tϕq - θsP::Tθ + ϕqIc::TI + θPm::Tθ dϕmvecs::TD - dθsP::TP hw_channel::THWC corendsM::TC transM::TM # need to store KEYS in type parameter otherwise allocation in NamedTuple(keys) keys_nljoint_args_indi::Val{KEYS} nljoint_args_fix::TAF - n_MC_val::Val{NMC} end function (f::ForwardDiffGradPointlossiZCl{KEYS})(tup) where KEYS local hi, nljoint_args_indi_tup, ϕm, i @@ -127,7 +122,7 @@ function (f::ForwardDiffGradPointlossiZCl{KEYS})(tup) where KEYS #nljoint_args_indi = NamedTuple{f.keys_nljoint_args_indi}(nljoint_args_indi_tup) nljoint_args_indi = NamedTuple{KEYS}(nljoint_args_indi_tup) forwarddiff_grad_pointloss_z!(hi, f.approx, nljoint_args_indi, ϕm, i, - f.ϕqIc, f.θsP, f.dϕmvecs, f.dθsP, f.hw_channel, f.corendsM, f.transM, + f.ϕqIc, f.θPm, f.dϕmvecs, f.hw_channel, f.corendsM, f.transM, f.nljoint_args_fix, nothing, #true, @@ -138,8 +133,8 @@ end # ϕm vector version # TOOD pass closure object rather than many parameters function forwarddiff_grad_pointloss_z!(hi, approx::AbstractHVIApproximation, - nljoint_args_indi, ϕm::AbstractVector{TF}, i, - ϕqIc, θsP, dϕmvecs, dθsP, hw_channel::Channel, cor_endsM, transM::Stacked, + nljoint_args_indi, ϕm::AbstractMatrix{TF}, i::Integer, + ϕqIc::AbstractVector, θP::AbstractMatrix, dϕmvecs, hw_channel::Channel, cor_endsM, transM::Stacked, nljoint_args_fix, omit_gradient=nothing) where {TF} # need to calls to ForwardDiff.gradient! to avoid reshape @@ -149,397 +144,119 @@ function forwarddiff_grad_pointloss_z!(hi, approx::AbstractHVIApproximation, with_channel_element(hw_channel) do hwi_tup hwi, hwi_dc = hwi_tup #grad_conf = hwi.grad_conf - nelboi_z_mat, nelboi_z_vec = make_pointlossz_lazy_cls(hi, hwi_dc, approx, nljoint_args_indi, - ϕm, ϕqIc, θsP, + pointlossi_z_vec = make_pointlossz_lazy_cls(hi, nljoint_args_indi, cor_endsM, transM, nljoint_args_fix, - ax_ϕqI, Val(1), + ax_ϕqI, ) - grad_θsP = if isnothing(omit_gradient) - ForwardDiff.gradient!(hwi.grad_θsP, nelboi_z_mat, CA.getdata(θsP), hwi.grad_conf_mat) - hwi.grad_θsP + grad_θP, grad_ϕm, grad_ϕqI = if isnothing(omit_gradient) + local θP_ϕm_ϕqI + # transform ϕqI to a 1-column matrix to vcat it to θP and ϕm + ϕqI = if strides(CA.getdata(ϕqIc))[1] == 1 + # to avoid allocations with reshape + Base.ReshapedArray(CA.getdata(ϕqIc), (length(CA.getdata(ϕqIc)), 1), ()) + else + reshape(CA.getdata(ϕqIc), :, 1) + end + () -> begin + tmpf(ϕqI) = mapreduce(identity, +, view(ϕqI,:,1)) + @allocated tmpf(ϕqI) # zero + end + θP_ϕm_ϕqI = Vcat(θP, ϕm, ϕqI) + tmp = similar(θP_ϕm_ϕqI) + ForwardDiff.gradient!(hwi.grad_θsP_ϕm_ϕqI, pointlossi_z_vec, θP_ϕm_ϕqI, hwi.grad_conf_point) + hwi.grad_θsP_ϕm_ϕqI.args[1], hwi.grad_θsP_ϕm_ϕqI.args[2], hwi.grad_θsP_ϕm_ϕqI.args[3] else - θsP - end - grad_ϕm, grad_ϕqI = if isnothing(omit_gradient) - local ϕm_ϕqI - ϕm_ϕqI = Vcat(ϕm, CA.getdata(ϕqIc)) - ForwardDiff.gradient!(hwi.grad_ϕm_ϕqI, nelboi_z_vec, ϕm_ϕqI, hwi.grad_conf_vec) - hwi.grad_ϕm_ϕqI.args[1], hwi.grad_ϕm_ϕqI.args[2] - else - ϕm, CA.getdata(ϕqIc) + θP, ϕm, CA.getdata(ϕqIc) end copyto!(view(dϕmvecs, :, i), grad_ϕm) # returning SVector helps avoiding allocations during reduce local res res = (; dϕqI = SA.SVector{axis_length(CA.getaxes(ϕqIc)[1])}(grad_ϕqI), - dθsP = SA.SVector{axis_length(ax_ϕP)*1}(grad_θsP), - ) - end -end - - -# ϕm matrix version -function forwarddiff_grad_pointloss_z!(hi, approx::AbstractHVIApproximation, - nljoint_args_indi, ϕm::AbstractMatrix{TF}, i, - ϕqIc, θsP, dϕmvecs, dθsP, hw_channel::Channel, cor_endsM, transM::Stacked, - nljoint_args_fix, - omit_gradient=nothing) where {TF} - # need to calls to ForwardDiff.gradient! to avoid reshape - # group arrays of same dimensionality (here grad_ϕm and grad_ϕqI) into a Vcat - ax_ϕqI = CA.getaxes(ϕqIc)[1] - ax_ϕqP = get_axis_θP(nljoint_args_fix.f!) - with_channel_element(hw_channel) do hwi_tup - hwi, hwi_dc = hwi_tup - #grad_conf = hwi.grad_conf - nelboi_z_mat, nelboi_z_vec = make_pointlossz_lazy_cls(hi, hwi_dc, approx, nljoint_args_indi, - ϕm, ϕqIc, θsP, - cor_endsM, transM, nljoint_args_fix, - ax_ϕqI, Val(2), ## 2-dimensional ϕm - ) - grad_θsP, grad_ϕm = if isnothing(omit_gradient) - local θsP_ϕm - θsP_ϕm = Vcat(θsP, ϕm) - ForwardDiff.gradient!(hwi.grad_θsP_ϕm, nelboi_z_mat, θsP_ϕm, hwi.grad_conf_mat) - hwi.grad_θsP_ϕm.args[1], hwi.grad_θsP_ϕm.args[2] - else - θsP, ϕm - end - grad_ϕqI = if isnothing(omit_gradient) - # TODO try reusing a view into grad_ϕm_ϕqI - ForwardDiff.gradient!(hwi.grad_ϕqI, nelboi_z_vec, CA.getdata(ϕqIc), hwi.grad_conf_vec) - hwi.grad_ϕqI - else - CA.getdata(ϕqIc) - end - copyto!(view(dϕmvecs, :, i), grad_ϕm) - # cannot use SVector for - # dθsP .+= grad_θsP # SharedArray, - # # returning SVector helps avoiding allocations during reduce - # local dϕqIc - # dϕqIc = SA.SVector{axis_length(CA.getaxes(ϕqIc)[1])}(grad_ϕqI) - local res - res = (; - dϕqI = SA.SVector{axis_length(CA.getaxes(ϕqIc)[1])}(grad_ϕqI), - dθsP = SA.SVector{axis_length(ax_ϕqP)*1}(grad_θsP), + dθsP = SA.SVector{axis_length(ax_ϕP)*1}(grad_θP), ) end end - -function make_pointlossz_lazy_cls(hi, hwi_dc, approx::AbstractHVIApproximation, nljoint_args_indi, - ϕm, ϕqIc, θsP, - cor_endsM, transM, nljoint_args_fix, - ax_ϕqI::CA.Axis, ndim_ϕ::Val{1}, - ) - # mat, here holds only θsP - # vec, is a Vcat(ϕqI, ϕm) - # no views involved here - function nelboiz_cl_mat(mat::AbstractMatrix) - local θsP_dual, ϕm_dual, ϕqIc_dual - θsP_dual = mat - ϕm_dual = copyto!(PAT.get_tmp(hwi_dc.ϕm_dc, mat), ϕm) - ϕqIc_dual = CA.ComponentArray( - copyto!(PAT.get_tmp(hwi_dc.ϕqI_dc, mat), ϕqIc), CA.getaxes(ϕqIc)) - compute_pointlossi_z!( - hi, nljoint_args_indi, - ϕm_dual, ϕqIc_dual, θsP_dual, - cor_endsM, transM, nljoint_args_fix, - )[1] - end - function nelboiz_cl_vec(vec::Vcat) - local θsP_dual, ϕm_dual, ϕqIc_dual - ϕm_dual = vec.args[1] - ϕqIc_dual = CA.ComponentVector(vec.args[2], ax_ϕqI) - θsP_dual = copyto!(PAT.get_tmp(hwi_dc.θsP_dc, vec), θsP) # reshape(reinterpred( - compute_pointlossi_z!( - hi, nljoint_args_indi, - ϕm_dual, ϕqIc_dual, θsP_dual, - cor_endsM, transM, nljoint_args_fix, - )[1] - end - nelboiz_cl_mat, nelboiz_cl_vec -end - -# ϕm Matrix case -function make_pointlossz_lazy_cls(hi, hwi_dc, approx::AbstractHVIApproximation, rnormM, nljoint_args_indi, - ϕm, ϕqIc, θsP, +function make_pointlossz_lazy_cls(hi, nljoint_args_indi, cor_endsM, transM, nljoint_args_fix, - ax_ϕqI::CA.Axis, ndim_ϕ::Val{2}, + ax_ϕqI::CA.Axis, ) - # mat is Vcat(θsP, ϕm) - # vec is ϕqI - function nelboiz_cl_mat(mat::Vcat) - local θsP_dual, ϕm_dual, ϕqIc_dual - θsP_dual = mat.args[1] + # mat, is a Vcat(θP, ϕm, ϕqI) + function pointlossi_cl_vec(mat::Vcat) + local θP_dual, ϕm_dual, ϕqIc_dual + θP_dual = mat.args[1] ϕm_dual = mat.args[2] - ϕqIc_dual = CA.ComponentArray( - copyto!(PAT.get_tmp(hwi_dc.ϕqI_dc, mat), ϕqIc), CA.getaxes(ϕqIc)) - compute_pointlossi_z!( - hi, nljoint_args_indi, - ϕm_dual, ϕqIc_dual, θsP_dual, - cor_endsM, transM, nljoint_args_fix, - )[1] - end - function nelboiz_cl_vec(vec::AbstractVector) - local θsP_dual, ϕm_dual, ϕqIc_dual - ϕqIc_dual = CA.ComponentVector(vec, ax_ϕqI) - ϕm_dual = copyto!(PAT.get_tmp(hwi_dc.ϕm_dc, vec), ϕm) - θsP_dual = copyto!(PAT.get_tmp(hwi_dc.θsP_dc, vec), θsP) + #ϕqIc_dual = CA.ComponentVector(vec(mat.args[3]), ax_ϕqI) + ϕqIc_dual = CA.ComponentVector(view(mat.args[3],:,1), ax_ϕqI) compute_pointlossi_z!( hi, nljoint_args_indi, - ϕm_dual, ϕqIc_dual, θsP_dual, + ϕm_dual, ϕqIc_dual, θP_dual, cor_endsM, transM, nljoint_args_fix, )[1] end - nelboiz_cl_mat, nelboiz_cl_vec -end - - -function get_pullback_cl_sample_ζsP(::AbstractArray{TF}; n_θP, n_MC, sample_buffers) where TF - #dϕqc, dζsP, dlogσ_ζP, ζsP, logσ_ζP, rnormP, ϕqc) - ζsP_ = Matrix{TF}(undef, n_θP, n_MC) - logσ_ζP_ = Vector{TF}(undef, n_θP) - dζsP_ = similar(ζsP_) # allocate space for derivatives - dlogσ_ζP_ = similar(logσ_ζP_) - drnormP = similar(ζsP_) - dsample_buffers = map_leaves_nt(similar, sample_buffers) # allocate space for derivatives - # - function pullback_cl_sample_ζsP!(dϕqc, dζsP, dlogσ_ζP, ζsP, logσ_ζP, - approx::AbstractHVIApproximation, rnormP, ϕqc, cor_endsP, sample_buffers::NamedTuple) - Enzyme.make_zero!(dϕqc) - Enzyme.make_zero!(dsample_buffers) - fill!(dϕqc, 0) # the derivative to compute - copyto!(dζsP_, dζsP) # input cotangents (modified in-place by Enzyme) - copyto!(dlogσ_ζP_, dlogσ_ζP) - copyto!(ζsP_, ζsP) # modified in-place by Enzyme - copyto!(logσ_ζP_, logσ_ζP) - EnzBuffers = isempty(sample_buffers) ? Enzyme.Const(sample_buffers) : - Enzyme.Duplicated(sample_buffers, dsample_buffers) - Enzyme.autodiff( - Enzyme.Reverse, - sample_ζsP!, - Enzyme.Duplicated(ζsP_, dζsP_), - Enzyme.Duplicated(logσ_ζP_, dlogσ_ζP_), - Enzyme.Const(approx), - Enzyme.DuplicatedNoNeed(rnormP, drnormP), - Enzyme.Duplicated(ϕqc, dϕqc), - Enzyme.Const(cor_endsP), - EnzBuffers, - ) - end -end - -function get_pullback_g_apply(::AbstractArray{TG}, ::AbstractArray{TF}; - n_θP, n_cov, n_covP, n_MC, n_site, n_M, - ) where {TG, TF} - xMP_ = Matrix{TG}(undef, (n_cov + n_covP), n_MC * n_site) - dxMP_ = similar(xMP_) - ϕms_ = Matrix{TF}(undef, n_M, n_site) - ϕms_mcs_ = Array{TF,3}(undef, n_M, n_MC, n_site) - dϕms_ = similar(ϕms_) - dϕms_mcs_ = similar(ϕms_mcs_) - #dζsP_ = Matrix{TF}(undef, n_θP, n_MC) - # - function pullback_g_apply!(dϕg, dζsP, ϕms, dϕm, ϕg, xM, ζsP, - pbm_covar_indices, g, is_testmode) - ϕms_buffer = isnothing(pbm_covar_indices) ? ϕms_ : ϕms_mcs_ - dϕms_buffer = isnothing(pbm_covar_indices) ? dϕms_ : dϕms_mcs_ - # assert that buffers were constructed with correct sizes - n_cov_f, n_site_f = size(xM) - n_covP_f = isnothing(pbm_covar_indices) ? 0 : length(pbm_covar_indices) - n_MC_f = size(ζsP,2) - @assert (n_cov, n_covP, n_MC, n_site) == (n_cov_f, n_covP_f, n_MC_f, n_site_f) - @assert size(dϕms_buffer) == size(dϕm) - - fill!(dϕg, zero(eltype(dϕg))) - fill!(dζsP, zero(eltype(dζsP))) # also output cotangent - fill!(dxMP_, zero(eltype(dxMP_))) - # primal will be updated as in the forward, but shadow needs to be preserved - copyto!(ϕms_buffer, ϕms) # copy to avoid modifying ϕm (although should be the same) - copyto!(dϕms_buffer, dϕm) # copy to avoid modifying dϕm - #copyto!(dζsP_, dζsP) # output, does not need to be preserved - # - Enzyme.autodiff( - Enzyme.Reverse, - #does not make SimpleChains work - currently cannot use SimpleChains - #Enzyme.set_runtime_activity(Enzyme.Reverse), # TODO only activate for SimpleChains - g_apply!, - Enzyme.Duplicated(ϕms_buffer, dϕms_buffer), - Enzyme.Duplicated(ϕg, dϕg), - Enzyme.Const(xM), - Enzyme.Duplicated(ζsP, dζsP), - Enzyme.Const(pbm_covar_indices), - Enzyme.Const(g), - Enzyme.Duplicated(xMP_, dxMP_), - Enzyme.Const(is_testmode), - ) - end -end - -# two return values: primal and derivative -# given coderiv dy, parameters and helpers -# dϕm -> dϕg, dζsP -function get_pullback_g_apply_ϕms_mcs2D_buffer(::AbstractArray{TG}, ::AbstractArray{TF}; - n_θP, n_cov, n_covP, n_MC, n_site, n_M, - ) where {TG, TF} - xMP_ = Matrix{TG}(undef, (n_cov + n_covP), n_MC * n_site) - dxMP_ = similar(xMP_) - ϕms_ = Matrix{TF}(undef, n_M, n_site) - ϕms_mcs_ = Array{TF,3}(undef, n_M, n_MC, n_site) - dϕms_ = similar(ϕms_) - dϕms_mcs_ = similar(ϕms_mcs_) - #dζsP_ = Matrix{TF}(undef, n_θP, n_MC) - ϕms_mcs2D_buffer_ = Matrix{TF}(undef, n_M, n_MC * n_site) - dϕms_mcs2D_buffer_ = similar(ϕms_mcs2D_buffer_) - # - function pullback_g_apply_ϕms_mcs2D_buffer!(dϕg, dζsP, ϕms, dϕm, ϕg, xM, ζsP, - pbm_covar_indices, g, is_testmode, ϕms_mcs2D_buffer) - ϕms_buffer = isnothing(pbm_covar_indices) ? ϕms_ : ϕms_mcs_ - dϕms_buffer = isnothing(pbm_covar_indices) ? dϕms_ : dϕms_mcs_ - # assert that buffers were constructed with correct sizes - n_cov_f, n_site_f = size(xM) - n_covP_f = isnothing(pbm_covar_indices) ? 0 : length(pbm_covar_indices) - n_MC_f = size(ζsP,1) - @assert (n_cov, n_covP, n_MC, n_site) == (n_cov_f, n_covP_f, n_MC_f, n_site_f) - @assert size(dϕms_buffer) == size(dϕm) - @assert size(ϕms_mcs2D_buffer_) == size(ϕms_mcs2D_buffer) - - fill!(dϕg, zero(eltype(dϕg))) - fill!(dζsP, zero(eltype(dζsP))) # also output cotangent - fill!(dxMP_, zero(eltype(dxMP_))) - fill!(dϕms_mcs2D_buffer_, zero(eltype(dϕms_mcs2D_buffer_))) - # primal will be updated as in the forward, but shadow needs to be preserved - copyto!(ϕms_buffer, ϕms) # copy to avoid modifying ϕm (although should be the same) - copyto!(dϕms_buffer, dϕm) # copy to avoid modifying dϕm - #copyto!(dζsP_, dζsP) # output, does not need to be preserved - copyto!(ϕms_mcs2D_buffer_, ϕms_mcs2D_buffer) # copy to avoid modifying dϕm - # - Enzyme.autodiff( - Enzyme.Reverse, g_apply!, - Enzyme.Duplicated(ϕms_buffer, dϕms_buffer), - Enzyme.Duplicated(ϕg, dϕg), - Enzyme.Const(xM), - Enzyme.Duplicated(ζsP, dζsP), - Enzyme.Const(pbm_covar_indices), - Enzyme.Const(g), - Enzyme.Duplicated(xMP_, dxMP_), - Enzyme.Const(is_testmode), - Enzyme.Duplicated(ϕms_mcs2D_buffer_, dϕms_mcs2D_buffer_), - ) - end + pointlossi_cl_vec end -function get_pullback_cl_transformζ!(::AbstractArray{TF}; n_θ, n_MC) where {TF} - θs_buffer = Matrix{TF}(undef, n_θ, n_MC) - dθs_buffer = similar(θs_buffer) - # Enzyme seeds an Active return value with one, so to seed its cotangent - # with an arbitrary dladJacTP we fold it in as a scalar factor, relying on - # linearity of the reverse-mode adjoint: the pullback then delivers - # dladJacTP * ∂ladJacT/∂(·) to θs and ζs, as with the former Ref seeding. - function pullback_cl_transformζ!(dζs, dθs, dladJacTP, θs, trans, ζs) - fill!(dζs, zero(eltype(dζs))) - copyto!(θs_buffer, θs) - copyto!(dθs_buffer, dθs) - # Trick of seeding the active return value different to unity: - # Here dladJacTP is captured from the closure (an Active-compatible scalar), - # the returned Active value's unit seed gets multiplied by dladJacTP, - # and the adjoints of θs/ζs come out as dladJacTP * ∂/∂(·) - Enzyme.autodiff( - Enzyme.Reverse, - (θs_, trans_, ζs_) -> dladJacTP * transformζ!(θs_, trans_, ζs_), - Enzyme.Duplicated(θs_buffer, dθs_buffer), - Enzyme.Const(trans), - Enzyme.Duplicated(ζs, dζs), - ) - end -end -function prepare_gradelbo_helpers( - ϕqIc, ϕm::AbstractArray{TF,ND}, θsP, +function prepare_gradpointloss_helpers( + ϕqIc, ϕm::AbstractArray{TF,ND}, θsP::AbstractMatrix, ϕg::AbstractVector{TG}, ϕqPc::AbstractVector{TF}, approx::AbstractHVIApproximation; pbm_covar_indices, n_workers, - h, rnormMi, nljoint_args_indi, - n_site, n_cov, cor_ends, + h, nljoint_args_indi, + n_site, n_M, n_cov, cor_ends, transP::Stacked, transM::Stacked, nljoint_args_fix, - n_MC_val::Val{NMC}, diffchunk::ForwardDiff.Chunk{chunk} = ForwardDiff.Chunk(8), - ) where {TG, TF, ND, NMC, chunk} + ) where {TG, TF, ND, chunk} hi1 = h.helpers_sites[1] ax_θP = get_axis_θP(nljoint_args_fix.f!) ax_θM = get_axis_θM(nljoint_args_fix.f!) # TODO get n_site, n_cov from - cv_grad = CA.ComponentArray(; - ϕqIc, - ϕm, # = selectdim(ϕms, ndims(ϕm), 1), - θsP, - ) - # ϕqIc = cv_grad[Val(:ϕqIc)] - # ϕm = cv_grad[Val(:ϕm)] - # θsP = cv_grad[Val(:θsP)] - use_ϕm_vector = ND==1 # single ϕm across MC (for each site) n_covP = isnothing(pbm_covar_indices) ? 0 : length(pbm_covar_indices) n_ϕmvec = length(ϕm) #use_ϕm_matrix ? n_M : n_M * n_MC - n_θP, n_MC = size(θsP) - @assert n_MC == NMC + (n_θP,n_MC) = size(θsP) + @assert n_MC == 1 @assert n_θP == axis_length(ax_θP) - n_M = size(ϕm,1) + # n_M = size(ϕm,1) # To sync across procs/threads, use a Channel # https://juliafolds2.github.io/OhMyThreads.jl/stable/literate/tls/tls/#The-safe-way:-Channel # @assert length(nljoint_args_indi.xP) > 1 - # nelboi_z = make_nelboiz_cl(hi1, approx, rnormMi, nljoint_args_indi, CA.getaxes(cv_grad), + # pointlossi_z = make_pointlossiz_cl(hi1, approx, rnormMi, nljoint_args_indi, CA.getaxes(cv_grad), # cor_ends.M, transM, nljoint_args_fix) hw_dcs_raw = (; ϕm_dc = similar(ϕm), # callocation on using ComponentArray (_restructure), need to use plain vector ϕqI_dc = similar(CA.getdata(ϕqIc)), - θsP_dc = similar(θsP), + θP_dc = similar(θsP), ) # make sure, each worker gets its own copy, use similar + # TODO remove hw_dcs, not needed because Vcat of all vectors hw_dcs = (hw_dcs_raw) -> map_leaves_nt(x -> PAT.DiffCache(similar(x), chunk), hw_dcs_raw) - nelboi_z_mat, nelboi_z_vec = make_pointlossz_lazy_cls(hi1, hw_dcs(hw_dcs_raw), approx, rnormMi, nljoint_args_indi, - ϕm, ϕqIc, θsP, + pointlossi_z_vec = make_pointlossz_lazy_cls(hi1, nljoint_args_indi, cor_ends.M, transM, nljoint_args_fix, - CA.getaxes(ϕqIc)[1], Val(ND), + CA.getaxes(ϕqIc)[1], ) - grad_conf_mat, grad_conf_vec = if use_ϕm_vector - local grad_ϕm_ϕqI1, diffchunks_vec, diffchunks_mat - diffchunk_vec = ForwardDiff.Chunk(min(chunk, - axis_length(CA.getaxes(ϕqIc)[1]) + axis_length(ax_θM))) - diffchunk_mat = ForwardDiff.Chunk(min(chunk, - axis_length(CA.getaxes(ϕqPc)[1]))) - grad_conf_mat_ϕmvec() = ForwardDiff.GradientConfig(nelboi_z_mat, CA.getdata(θsP), diffchunk_mat) - grad_ϕm_ϕqI1 = grad_ϕm_ϕqI(ϕm, ϕqIc) - function grad_conf_vec_ϕmvec() - cfg_vec = ForwardDiff.GradientConfig(nelboi_z_vec, grad_ϕm_ϕqI1, diffchunk_vec) - update_duals_Vcat(cfg_vec, grad_ϕm_ϕqI1) - end - grad_conf_mat_ϕmvec, grad_conf_vec_ϕmvec - else - local grad_θP_ϕm1, diffchunks_vec, diffchunks_mat - diffchunks_vec = ForwardDiff.Chunk(min(chunk, - axis_length(CA.getaxes(ϕqIc)[1]))) - diffchunks_mat = ForwardDiff.Chunk(min(chunk, - NMC*(axis_length(CA.getaxes(ϕqPc)[1]) + axis_length(ax_θM)))) - grad_conf_vec_ϕmmat() = ForwardDiff.GradientConfig(nelboi_z_vec, CA.getdata(ϕqIc), diffchunks_vec) - grad_θP_ϕm1 = grad_θsP_ϕm(ϕm, θsP) - function grad_conf_mat_ϕmmat() - cfg_mat = ForwardDiff.GradientConfig(nelboi_z_mat, grad_θP_ϕm1, diffchunks_mat) - update_duals_Vcat(cfg_mat, grad_θP_ϕm1) - end - grad_conf_mat_ϕmmat, grad_conf_vec_ϕmmat + local grad_θsP_ϕm_ϕqI1, diffchunks_point + ϕqIm = reshape(CA.getdata(ϕqIc),:,1) + ϕm_mat = (ND == 2) ? ϕm : reshape(ϕm, :, 1) + grad_θsP_ϕm_ϕqI1 = grad_θsP_ϕm_ϕqI(θsP, ϕm_mat, ϕqIm) + # n_M is currently not static, hence diffchunks_vec and gradientConfig not inferred + # however, n_θM is a lower bound, so use this one to set a minimum length + # chunk size might be not as large as possible, but safely not larger than paramvector + diffchunks_point = ForwardDiff.Chunk(min(chunk, + #axis_length(ax_θP) + n_M + axis_length(CA.getaxes(ϕqIc)[1]))) + axis_length(ax_θP) + axis_length(ax_θM) + axis_length(CA.getaxes(ϕqIc)[1]))) + function grad_conf_point() + cfg_vec = ForwardDiff.GradientConfig(pointlossi_z_vec, grad_θsP_ϕm_ϕqI1, diffchunks_point) + update_duals_Vcat3(cfg_vec, grad_θsP_ϕm_ϕqI1) end get_helpers_worker = () -> begin (; - grad_conf_vec = grad_conf_vec(), - grad_conf_mat = grad_conf_mat(), - grad_θsP = similar(θsP), - grad_ϕm_ϕqI = grad_ϕm_ϕqI(ϕm, ϕqIc), - grad_θsP_ϕm= grad_θsP_ϕm(ϕm, θsP), - grad_ϕqI = similar(CA.getdata(ϕqIc)), + grad_conf_point = grad_conf_point(), + grad_θsP_ϕm_ϕqI = grad_θsP_ϕm_ϕqI(θsP, ϕm_mat, ϕqIm), ) end h1 = get_helpers_worker() @@ -555,39 +272,19 @@ function prepare_gradelbo_helpers( dϕg = similar(ϕg), dϕqPc = similar(ϕqPc), dϕqIc = similar(ϕqIc), - dθsP = similar(θsP), ∂elbo_∂θP = similar(θsP), ∂elbo_∂ϕm_∂ζP = similar(θsP), ∂elbo_∂logσ_ζP = Vector{TF}(undef, size(θsP,1)), ∂elbo_∂θP_∂ζP = Matrix{TF}(undef, size(θsP)), ∂elbo_∂ϕqm = Array{TF}(undef, size(ϕm)..., n_site), hw_channel, - pullback_cl_sample_ζsP! = get_pullback_cl_sample_ζsP(CA.getdata(ϕqPc); - n_θP, n_MC, sample_buffers = h.sample_buffers), pullback_cl_transformζsP! = get_pullback_cl_transformζ!(CA.getdata(ϕqPc); n_θ = n_θP, n_MC), pullback_g_apply! = get_pullback_g_apply( ϕg, ϕqPc; n_θP, n_cov, n_covP, n_MC, n_site, n_M), # SharedArrays last because of problems with ReferenceRevision dϕmvecs = SharedArrays.SharedArray{TF}(n_ϕmvec, n_site), - #dθsP = SharedArrays.SharedArray{TF}(n_θP, n_MC), ) end -@inline grad_ϕm_ϕqI(ϕm::AbstractVector, ϕqIc) = Vcat(similar(ϕm), similar(CA.getdata(ϕqIc))) -@inline grad_ϕm_ϕqI(ϕm::AbstractMatrix, ϕqIc) = Vcat(similar(ϕm[1:0,:1]), similar(CA.getdata(ϕqIc))) # not used but preallocated -@inline grad_θsP_ϕm(ϕm::AbstractVector, θsP) = Vcat(similar(θsP[:,1:0]), similar(θsP[:,1:0])) # not used but preallocated -@inline grad_θsP_ϕm(ϕm::AbstractMatrix, θsP) = Vcat(similar(θsP), similar(ϕm)) - - -@inline function update_duals_Vcat(cfg::ForwardDiff.GradientConfig{T,V,N}, vc::Vcat) where {T,V,N} - # similar Vcat returns a vector, need to convert dual of cfg to Vcat - # no view, but copy with index_at_dim - duals_vcat = Vcat(index_at_dim(cfg.duals, axes(vc.args[1], 1); dim=1), - index_at_dim(cfg.duals, axes(vc.args[2], 1); dim=1)); - cfg = ForwardDiff.GradientConfig{T,V,N,typeof(duals_vcat)}(cfg.seeds, duals_vcat) - @assert cfg.duals isa Vcat - cfg -end - - - +# need similar to make a copy for each worker +grad_θsP_ϕm_ϕqI(θP, ϕm, ϕqIc) = Vcat(similar(θP), similar(ϕm), similar(CA.getdata(ϕqIc))) diff --git a/test/runtests.jl b/test/runtests.jl index b1bc06a..517ee4f 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -56,6 +56,7 @@ const GROUP = get(ENV, "GROUP", "All") # defined in in CI.yml @time @safetestset "BMApplicatorMC" include("test_PBMApplicatorMC.jl") @time @safetestset "util" include("test_util.jl") @time @safetestset "elbo_sites" include("test_elbo_site.jl") + @safetestset "elbo_sites_doubleMM" include("test_elbo_site_doubleMM.jl") end end diff --git a/test/test_elbo_site.jl b/test/test_elbo_site.jl index 25949b6..0c11d18 100644 --- a/test/test_elbo_site.jl +++ b/test/test_elbo_site.jl @@ -13,9 +13,9 @@ using LinearAlgebra using StatsFuns: logistic using Test -Distributed.@everywhere using HybridVariationalInference using HybridVariationalInference: HybridVariationalInference as CP using HybridVariationalInference: HybridVariationalInference as HVI +Distributed.@everywhere using HybridVariationalInference using StableRNGs using Random using ComponentArrays: ComponentArrays as CA diff --git a/test/test_elbo_site_doubleMM.jl b/test/test_elbo_site_doubleMM.jl index 7e36216..d4a67d2 100644 --- a/test/test_elbo_site_doubleMM.jl +++ b/test/test_elbo_site_doubleMM.jl @@ -155,16 +155,40 @@ end nothing end # allocations in StatefulLuxLayer - #@profview_allocs loop_pointloss_sites!(100_000, hp, ϕ, sample_args, indiv_args, nljoint_args_fix) + #@profview_allocs loop_pointloss_sites!(10_000, hp, ϕ, sample_args, indiv_args, nljoint_args_fix) function alloc_pointloss_sites!(hp, ϕ, sample_args, indiv_args, nljoint_args_fix) + loop_pointloss_sites!(1, hp, ϕ, sample_args, indiv_args, nljoint_args_fix) @allocated loop_pointloss_sites!(10, hp, ϕ, sample_args, indiv_args, nljoint_args_fix) end + alloc_pointloss_sites!(hp, ϕ, sample_args, indiv_args, nljoint_args_fix) alloc_pointloss_sites!(hp, ϕ, sample_args, indiv_args, nljoint_args_fix) <= 47_200 # - # res0, gradh0 = CP.grad_pointloss_sites( - # #@descend_code_warntype CP.neg_elbo_sites!( - # hp, (;), ϕ, sample_args, indiv_args, nljoint_args_fix) - + res0, gradh0 = CP.grad_pointloss_sites( + #@descend_code_warntype CP.pointloss_sites!( + hp, (;), ϕ, sample_args, indiv_args, nljoint_args_fix) + @test res0.dϕqPc.μζP[1] != 0 + @test res0.dϕqPc.logσ_ζP[1] == 0 + @test res0.dϕqIc.logσ_ζM_offsets[1] == 0 + function loop_grad_pointloss_sites( + n, hp, gradh0, ϕ, sample_args, indiv_args, nljoint_args_fix) + for _ in 1:n + @noinline CP.grad_pointloss_sites( + hp, gradh0, ϕ, sample_args, indiv_args, nljoint_args_fix) + end + nothing + end + function alloc_grad_pointloss_sites(hp, gradh0, ϕ, sample_args, indiv_args, nljoint_args_fix) + # avoid global variables -> pass them through function + loop_grad_pointloss_sites(1, hp, gradh0, ϕ, sample_args, indiv_args, nljoint_args_fix) + # fewer allocations than calls + #@profview_allocs loop_grad_pointloss_sites(1_000,hp, gradh0, ϕ, sample_args, indiv_args, nljoint_args_fix) + alloc = @allocated loop_grad_pointloss_sites(100,hp, gradh0, ϕ, sample_args, indiv_args, nljoint_args_fix) + @test alloc <= 3_011_200 + alloc + end + alloc_grad_pointloss_sites(hp, gradh0, ϕ, sample_args, indiv_args, nljoint_args_fix) + + end