Skip to content

implement point gradient #43

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
bgctw merged 4 commits into from
Oct 26, 2026
Merged

implement point gradient #43

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
3a1965f
trhead AbstractElboLossMode
bgctw Oct 24, 2026
26f00e3
implement set_ζsP! and set_ζsM!
bgctw Oct 24, 2026
ac0ecef
implement grad_neg_elbo_sites for NegLogDensity()
bgctw Oct 26, 2026
f51bd8d
test point gradient by Enzyme
bgctw Oct 26, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions src/HVIApproximation.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,8 @@ struct MeanHVIApproximationDev <: AbstractMeanHVIApproximation end

# neglecting correlations
struct DiagonalHVIApproximation <: AbstractHVIApproximation end
get_numberof_θM(::DiagonalHVIApproximation, ml_pred::AbstractArray) = length(ml_pred)



abstract type AbstractMeanVarSepHVIApproximation <: AbstractHVIApproximation end
Expand Down
50 changes: 44 additions & 6 deletions src/elbo_site.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,13 +6,21 @@
nothing
end

"""
Dispatch to make some part of neg_elbo_sites! polymorphic
"""
abstract type AbstractElboLossMode end
struct NegElbo <: AbstractElboLossMode end
struct NegLogDensity <: AbstractElboLossMode end

"""
elbo_helpers need to be initialized with new random numbers
in h.ζsP and hi.ζsM by calling randnPM! before.
By this way, we can compute the derivative corresponding to the forward pass
"""
function neg_elbo_sites!(
elbo_helpers::NamedTuple, # tuple of preallocated arrays
loss_mode::AbstractElboLossMode,
ϕ::NamedTuple,
rnormPM::NamedTuple,
sample_args::NamedTuple,
Expand All @@ -34,7 +42,7 @@
ϕqIc = intϕqI(ϕqI)
pbm_covar_indices = !isnothing(pbm_covar_indices) && isempty(pbm_covar_indices) ? nothing : pbm_covar_indices
#
sample_ζsP!(h.ζsP, h.logσ_ζP, approx, rnormPM.P, ϕqPc, cor_ends.P, h.sample_buffers) # n_P * n_MC
set_ζsP!(h.ζsP, h.logσ_ζP, loss_mode, approx, rnormPM.P, ϕqPc, cor_ends.P, h.sample_buffers) # n_P * n_MC
# (n_M x n_sit) or (n_M x n_MC x n_sit)
ϕms_buffer_key = isnothing(pbm_covar_indices) ? :ϕms : :ϕms_mcs
#g_apply!(h[ϕms_buffer_key], ϕg, xM, h.ζsP, pbm_covar_indices, g, h.xMP, is_testmode, h.ϕms_mcs2D_buffer)
Expand All @@ -49,10 +57,10 @@
function compute_nelboi_z_cl!(hi, rnormM, nljoint_args_indi_tup, ϕm)
#nljoint_args_indi = NamedTuple{keys_nljoint_args_indi}(nljoint_args_indi_tup)
nljoint_args_indi = name_tuple(nljoint_args_indi_tup, keys_nljoint_args_indi_val)
compute_nelboi_z!(hi, approx, rnormM, nljoint_args_indi, ϕm,
compute_nelboi_z!(hi, loss_mode, approx, rnormM, nljoint_args_indi, ϕm,
ϕqIc, θsP, cor_ends.M, transM, nljoint_args_fix)
end
#res_site = map(compute_nelboi_z!, h.helpers_sites, rnormPM.M, i_sites_train, ϕm_it)
#res_site = map(compute_nelboi_z!, loss_mode, h.helpers_sites, rnormPM.M, i_sites_train, ϕm_it)
#MAYBE: distributed mapreduce:
# https://docs.julialang.org/en/v1/stdlib/Distributed/#Distributed.@distributed
# https://github.com/SupaeroDataScience/DE/blob/main/notebooks/Introduction%20to%20MapReduce.ipynb
Expand All @@ -64,35 +72,65 @@
elbo = elbo_z - exp_ladJacTP - sum(h.logσ_ζP)
end

name_tuple(tup, ::Val{KEYS}) where KEYS = NamedTuple{KEYS}(tup)
@inline function set_ζsP!(
ζsP, logσ_ζP, ::NegElbo, approx, rnormP, ϕqPc, cor_endsP, sample_buffers)
sample_ζsP!(ζsP, logσ_ζP, approx, rnormP, ϕqPc, cor_endsP, sample_buffers) # n_P * n_MC
end
@inline function set_ζsP!(
ζsP, logσ_ζP::AbstractMatrix{T}, ::NegLogDensity, approx, rnormP, ϕqPc, cor_endsP, sample_buffers) where T
ζsP .= CA.getdata(view(ϕqPc,Val(:μζP)))
logσ_ζP .= zero(T) # omit loss component of entropy of distribution
end
@inline function set_ζsP!(
ζsP, logσ_ζP::AbstractVector{T}, ::NegLogDensity, approx, rnormP, ϕqPc, cor_endsP, sample_buffers) where T
ζsP .= CA.getdata(view(ϕqPc,Val(:μζP)))
logσ_ζP .= zero(T) # omit loss component of entropy of distribution
end


name_tuple(tup, ::Val{KEYS}) where KEYS = NamedTuple{KEYS}(tup)


function compute_nelboi_z!(hi, approx::AbstractHVIApproximation,
function compute_nelboi_z!(hi, loss_mode::AbstractElboLossMode, approx::AbstractHVIApproximation,
rnormM, nljoint_args_indi, ϕm::AbstractArray, ϕqIc::AbstractArray, θsP,
cor_endsM, transM, nljoint_args_fix, )
# DM is 1 (vector) if no population covariates - constant across MC
# 2 (Matrix) if g is driven by population covariates - differs across MC
# on update -> sync corresponding function within grad_neg_elbo_sites
if hi.ζsM isa PAT.DiffCache
# Gradiend called twice because of different dimensions need to pick the Dual

Check warning on line 101 in src/elbo_site.jl

View workflow job for this annotation

GitHub Actions / Spell Check with Typos 

"Gradiend" should be "Gradient".
template = eltype(ϕqIc) <: ForwardDiff.Dual ? ϕqIc : θsP
hi = map_leaves_nt(x -> PAT.get_tmp(x, template), hi)
end
#ζsM, logσ_ζM, rnorm, ϕqc::AbstractVector{T}, ϕm::AbstractMatrix, buffer_nθM::AbstractVector
sample_ζsM!(hi.ζsM, hi.logσ_ζM, approx, rnormM, ϕqIc, ϕm, cor_endsM, hi.sample_buffers)
set_ζsM!(hi.ζsM, hi.logσ_ζM, loss_mode, approx, rnormM, ϕqIc, ϕm, cor_endsM, hi.sample_buffers)
exp_ladJacTM = transformζ!(hi.θsM, transM, hi.ζsM) # return value captures ladJacT
# first component needs to be the full elbo
exp_nL = exp_nLi!(hi.y, θsP, hi.θsM, nljoint_args_indi, nljoint_args_fix, hi)[1]
nelboi_z = exp_nL - exp_ladJacTM - sum(hi.logσ_ζM)
end

@inline function set_ζsM!(ζsM, logσ_ζM, ::NegElbo, approx, rnormM, ϕqIc, ϕm, cor_endsM, sample_buffers)
sample_ζsM!(ζsM, logσ_ζM, approx, rnormM, ϕqIc, ϕm, cor_endsM, sample_buffers)
end
@inline function set_ζsM!(ζsM, logσ_ζM::AbstractArray{T}, ::NegLogDensity, approx, rnormM, ϕqIc, ϕm::AbstractVector, cor_endsM, sample_buffers) where T
ζsM .= view(ϕm, 1:size(ζsM,1)) # recycling fills all rows
logσ_ζM .= zero(T)
nothing
end
@inline function set_ζsM!(ζsM, logσ_ζM::AbstractArray{T}, ::NegLogDensity, approx, rnormM, ϕqIc, ϕm::AbstractMatrix, cor_endsM, sample_buffers) where T
ζsM .= view(ϕm, 1:size(ζsM,1), :) # μζM in first n_θM entries
logσ_ζM .= zero(T)
nothing
end

# get_tmp_rec_(x::PAT.DiffCache{<:AbstractArray}, template) = PAT.get_tmp(x, template)
# function get_tmp_rec_(x::Union{Tuple,NamedTuple}, template)
# map(xi -> get_tmp_rec_(xi, template), x)
# end



function prepare_rnorm(::AbstractVector{TF}; n_θP, n_θM, n_site, n_MC) where TF
(;
P = Matrix{TF}(undef, n_θP, n_MC),
Expand Down Expand Up @@ -344,12 +382,12 @@
(;xP, y_o, y_unc, i_indiv_train) = nljoint_args_indi
@assert size(y_o) == size(y_unc) == (n_obs,)
@assert size(θsP,2) == size(θsM,2) == n_MC
(;f!,fLogLik) = nljoint_args_fix

Check warning on line 385 in src/elbo_site.jl

View workflow job for this annotation

GitHub Actions / Spell Check with Typos 

"Lik" should be "Like" or "Lick" or "Link".
#θsPc = CA.ComponentMatrix(θsP, axθP, CA.FlatAxis())
#θsMc = CA.ComponentMatrix(θsM, axθM, CA.FlatAxis())
apply_model!(y_pred, f!, θsP, θsM, xP)
nL_reg = sum(y_pred) / n_MC
nL = fLogLik(y_pred, y_o, y_unc, hi.y_buffer) / n_MC

Check warning on line 390 in src/elbo_site.jl

View workflow job for this annotation

GitHub Actions / Spell Check with Typos 

"Lik" should be "Like" or "Lick" or "Link".
# nL = (5 * sum(θsP) + 3 * sum(θsM)) / n_MC
(; nL,)
# loss_comps = neg_elbo_ζtf(
Expand Down Expand Up @@ -399,7 +437,7 @@
pred
end

function fLogLik_dummy(y_pred, y_o, y_unc, y_buffer=similar(y_pred))

Check warning on line 440 in src/elbo_site.jl

View workflow job for this annotation

GitHub Actions / Spell Check with Typos 

"Lik" should be "Like" or "Lick" or "Link".
sum(y_pred)
end

66 changes: 52 additions & 14 deletions src/elbo_site_grad.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,12 +2,13 @@
function grad_neg_elbo_sites(
elbo_helpers::NamedTuple, # tuple of preallocated arrays
grad_elbo_helpers::NamedTuple, # tuple of preallocated arrays pullback closures
loss_mode,
ϕ::NamedTuple,
rnormPM::NamedTuple,
sample_args::NamedTuple,
indiv_args::NamedTuple,
nljoint_args_fix::NamedTuple;
n_MC_val::Val{NMC}, # need to pass by value for type inference on chunks
n_MC_val::Val{NMC_ELBO}, # need to pass by value for type inference on chunks
# approx::AbstractHVIApproximation,
# rnormPM::NamedTuple, # tuple of random numbers
# ϕg::AbstractVector{TG}, ϕqP::AbstractVector{TF}, ϕqI::AbstractVector{TF}, g,
Expand All @@ -21,7 +22,13 @@ function grad_neg_elbo_sites(
# is_testmode,
n_workers = 1, # TODO compute from executor and Distributed.nworkers and nthreads,...
executor::Transducers.Executor = Transducers.SequentialEx(),
) where NMC
) where NMC_ELBO
if loss_mode isa NegLogDensity
NMC = 1
n_MC_val = Val(1)
else
NMC = NMC_ELBO
end
h = elbo_helpers # preallocated μζP, dμζP, ζsP, ϕms, xMP, dxMP
(;ϕg, ϕqP, ϕqI) = ϕ
TF = eltype(ϕqP)
Expand All @@ -48,10 +55,11 @@ function grad_neg_elbo_sites(
ϕqIc, selectdim(h_ϕm, ndims(h_ϕm), 1), h.ζsP, ϕg, ϕqPc, approx;
pbm_covar_indices, n_workers,
h, rnormMi = rnormPM.M[1], nljoint_args_indi,
n_site, n_cov, cor_ends, transP, transM, nljoint_args_fix, n_MC_val)
n_site, n_cov, cor_ends, transP, transM, nljoint_args_fix, loss_mode,
n_MC_val)
hw_channel = gradh.hw_channel
#
sample_ζsP!(h.ζsP, h.logσ_ζP, approx, rnormPM.P, ϕqPc, cor_ends.P, h.sample_buffers) # n_P * n_MC
set_ζsP!(h.ζsP, h.logσ_ζP, loss_mode, approx, rnormPM.P, ϕqPc, cor_ends.P, h.sample_buffers) # n_P * n_MC
g_apply!(h_ϕm, ϕg, xM, h.ζsP, pbm_covar_indices, g, h.xMP, is_testmode)
ladJacTP = transformζ!(h.θsP, transP, h.ζsP) # return value captures ladJacT
#
Expand All @@ -67,7 +75,7 @@ function grad_neg_elbo_sites(
end
cl = ForwardDiffGradNelboiZCl(
approx, ϕqIc, gradh.dϕmvecs, gradh.hw_channel, cor_ends.M, transM,
Val(keys(nljoint_args_inds)), nljoint_args_fix, n_MC_val)
Val(keys(nljoint_args_inds)), nljoint_args_fix, loss_mode, n_MC_val)
ϕm_it = eachslice(h_ϕm; dims = ndims(h_ϕm))
# _a1 = Test.@inferred first(zip(h.helpers_sites, rnormPM.M, zip_eachlastdims(nljoint_args_inds), ϕm_it, 1:n_site))
# Test.@inferred cl(_a1)
Expand All @@ -87,7 +95,7 @@ function grad_neg_elbo_sites(
copyto!(∂elbo_∂θP, res.dθsP)
copyto!(gradh.dϕqIc, res.dϕqI)
# other gradients are one, calculated from formula
gradh.∂elbo_∂logσ_ζP .= -one(TF)
set_∂elbo_∂logσ_ζP!(gradh.∂elbo_∂logσ_ζP, loss_mode)
∂elbo_∂ladJacTP = -one(TF)
#
# pullback gradients of ϕqm -> gradh.dϕg and gradh.dζsP
Expand All @@ -107,18 +115,39 @@ function grad_neg_elbo_sites(
)
#
# pullback gradients of ∂elbo_∂ζP, ∂elbo_∂ϕm_∂ζP, and ∂elbo_∂logσ_ζP to dϕqP
gradh.pullback_cl_sample_ζsP!(
gradh.dϕqPc,
pullack_sample_ζsP!(gradh.dϕqPc, loss_mode,
∂elbo_∂θP_∂ζP + gradh.∂elbo_∂ϕm_∂ζP, gradh.∂elbo_∂logσ_ζP,
gradh,
h.ζsP, h.logσ_ζP, approx, rnormPM.P, ϕqPc, cor_ends.P, h.sample_buffers,
)
(;dϕqPc = gradh.dϕqPc, dϕqIc = gradh.dϕqIc, dϕg = gradh.dϕg), gradh
end

function set_∂elbo_∂logσ_ζP!(∂elbo_∂logσ_ζP::AbstractVector{T}, ::NegElbo) where T
∂elbo_∂logσ_ζP .= -one(T)
end
function set_∂elbo_∂logσ_ζP!(∂elbo_∂logσ_ζP::AbstractVector{T}, ::NegLogDensity) where T
∂elbo_∂logσ_ζP .= zero(T)
end


function pullack_sample_ζsP!(dϕqPc, ::NegElbo, dζsP, dlogσ_ζP, gradh, args...)
gradh.pullback_cl_sample_ζsP!(dϕqPc, dζsP, dlogσ_ζP, args...)
end
function pullack_sample_ζsP!(dϕqPc, ::NegLogDensity, dζsP, dlogσ_ζP::AbstractVector{TF}, args...) where TF
#error("debug pullack_sample_ζsP!(dϕqPc, ::NegLogDensity")
n_θM = size(dζsP, 1)
@assert all(dlogσ_ζP .== zero(TF))
dϕqPc[1:n_θM] .= dζsP
dϕqPc[(n_θM+1):end] .= zero(TF)
dϕqPc
end


"""
Callable to make deliver arguments that do not differ by individual to Foldl.mapreduce.
"""
struct ForwardDiffGradNelboiZCl{KEYS, NMC, TA, Tϕq, TD, THWC, TC, TM, TAF}
struct ForwardDiffGradNelboiZCl{KEYS, NMC, TA, Tϕq, TD, THWC, TC, TM, TAF, LM}
approx::TA
ϕqIc::Tϕq # TODO remove from closure, provided with preallocated VcatVecViews
dϕmvecs::TD
Expand All @@ -128,6 +157,7 @@ struct ForwardDiffGradNelboiZCl{KEYS, NMC, TA, Tϕq, TD, THWC, TC, TM, TAF}
# need to store KEYS in type parameter otherwise allocation in NamedTuple(keys)
keys_nljoint_args_indi::Val{KEYS}
nljoint_args_fix::TAF
loss_mode::LM
n_MC_val::Val{NMC}
end
function (f::ForwardDiffGradNelboiZCl{KEYS})(tup) where KEYS
Expand All @@ -152,7 +182,7 @@ function forwarddiff_grad_nelboi_z!(hi, rnormM,
hwi, hwi_dc = hwi_tup
#grad_conf = hwi.grad_conf
nelboi_vv = make_nelboiz_lazy_cls_vv(hi, cl.approx, rnormM, nljoint_args_indi,
cl.cor_endsM, cl.transM, cl.nljoint_args_fix, ax_ϕqI)
cl.cor_endsM, cl.transM, cl.nljoint_args_fix, cl.loss_mode, ax_ϕqI)
# grads are vector-shaped, but copyto! dϕmvecs and SAVector use those anyhow
grad_θsP, grad_ϕm, grad_ϕqI = if isnothing(omit_gradient)
# θsP_ϕm_ϕqI is preallocated per worker (hwi.θsP_ϕm_ϕqI).
Expand Down Expand Up @@ -184,6 +214,7 @@ end
function make_nelboiz_lazy_cls_vv(hi, approx::AbstractHVIApproximation,
rnormM, nljoint_args_indi,
cor_endsM, transM, nljoint_args_fix,
loss_mode::AbstractElboLossMode,
ax_ϕqI::CA.Axis
)
# VcatVecViews(θsP, ϕm, ϕqI)
Expand All @@ -194,7 +225,7 @@ function make_nelboiz_lazy_cls_vv(hi, approx::AbstractHVIApproximation,
ϕm_dual = subarray(vv,Val(2))
ϕqIc_dual = CA.ComponentVector(subarray(vv,Val(3)), ax_ϕqI)
compute_nelboi_z!(
hi, approx, rnormM, nljoint_args_indi,
hi, loss_mode, approx, rnormM, nljoint_args_indi,
ϕm_dual, ϕqIc_dual, θsP_dual,
cor_endsM, transM, nljoint_args_fix,
)[1]
Expand Down Expand Up @@ -369,9 +400,16 @@ function prepare_gradelbo_helpers(
n_site, n_cov, cor_ends,
transP::Stacked, transM::Stacked,
nljoint_args_fix,
n_MC_val::Val{NMC},
loss_mode::AbstractElboLossMode,
n_MC_val::Val{NMC_ELBO},
max_diffchunk::ForwardDiff.Chunk{chunk} = ForwardDiff.Chunk(8),
) where {TG, TF, ND, NMC, chunk}
) where {TG, TF, ND, NMC_ELBO, chunk}
if loss_mode isa NegLogDensity
NMC = 1
n_MC_val = Val(1)
else
NMC = NMC_ELBO
end
local hi1, ax_θP, ax_θM, n_covP, n_ϕmvec, n_θP, n_MC, n_M
hi1 = h.helpers_sites[1]
ax_θP = get_axis_θP(nljoint_args_fix.f!)
Expand Down Expand Up @@ -401,7 +439,7 @@ function prepare_gradelbo_helpers(
#
local nelboiz_vv, θsP_ϕm_ϕqI_vv1, diffchunks_elbo_vv
nelboiz_vv = make_nelboiz_lazy_cls_vv(hi1, approx, rnormMi, nljoint_args_indi,
cor_ends.M, transM, nljoint_args_fix, CA.getaxes(ϕqIc)[1], )
cor_ends.M, transM, nljoint_args_fix, loss_mode, CA.getaxes(ϕqIc)[1], )
θsP_ϕm_ϕqI_vv1 = θsP_ϕm_ϕqI_vv(θsP,ϕm,ϕqIc)
# n_M >= n_θM but not static
diffchunks_elbo_vv = ForwardDiff.Chunk(min(chunk,
Expand Down
Loading
Loading