Skip to content

Commit 748ae8f

Browse files
bgctwbgctw
committed
remove ZetaP from sample_ζresid_norm
1 parent bf58655 commit 748ae8f

5 files changed

Lines changed: 60 additions & 135 deletions

File tree

‎src/elbo.jl‎

Lines changed: 19 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -470,8 +470,9 @@ function generate_ζ(approx::AbstractMeanHVIApproximation, rng::AbstractRNG,
470470
# first pass: append μ_ζP_to covars, need ML prediction for magnitude of ζMs
471471
# TODO replace pbm_covar_indices by ComponentArray? dimensions to be type-inferred?
472472
xMP0 = _append_each_covars(xM, CA.getdata(μ_ζP), pbm_covar_indices)
473-
μ_ζMs0 = g(xMP0, ϕg; is_testmode)
474-
ζP_resids, ζMs_parfirst_resids, σ = sample_ζresid_norm(approx, rng, μ_ζP, μ_ζMs0, ϕq; n_MC, cor_ends, int_ϕq)
473+
ϕm0 = g(xMP0, ϕg; is_testmode)
474+
μ_ζMs0 = ϕm0
475+
ζP_resids, ζMs_parfirst_resids, σ = sample_ζresid_norm(approx, rng, ϕm0, ϕq; n_MC, cor_ends, int_ϕq)
475476
ζsP = isempty(μ_ζP) ? ζP_resids : (μ_ζP .+ ζP_resids) # n_par x n_MC
476477
if pbm_covar_indices isa SA.SVector{0}
477478
# do not need to predict again but just add the residuals to μ_ζP and μ_ζMs
@@ -488,7 +489,7 @@ function generate_ζ(approx::AbstractMeanHVIApproximation, rng::AbstractRNG,
488489
ζsMs_vec = map(eachcol(ζsP), eachslice(ζMs_parfirst_resids; dims=3)) do ζP, rMs
489490
# second pass: append ζP rather than μ_ζP to covars to xM
490491
xMP = _append_each_covars(xM, CA.getdata(ζP), pbm_covar_indices)
491-
μ_ζMst = g(xMP, ϕg)
492+
μ_ζMst = ϕm = g(xMP, ϕg; is_testmode)
492493
ζMs = (μ_ζMst .+ rMs)' # already transform to par-last form
493494
ζMs
494495
end
@@ -567,30 +568,36 @@ ML-model predcitions of size `(n_θM, n_site)`.
567568
ρsP, ρsM, logσ2_ζP, coef_logσ2_ζMs(intercept + slope),
568569
"""
569570
function sample_ζresid_norm(approx::AbstractHVIApproximation, rng::Random.AbstractRNG,
570-
ζP::AbstractVector, ζMs::AbstractMatrix, args...;
571+
ϕm::AbstractMatrix, ϕq::AbstractVector,
572+
args...;
571573
n_MC, cor_ends, int_ϕq)
574+
ζP = int_ϕq(CA.getdata(ϕq))[Val(:μP)]
575+
ζMs = ϕm
572576
n_θP, n_θMs = length(ζP), length(ζMs)
573577
# intm_PMs_parfirst = !isnothing(intm_PMs_parfirst) ? intm_PMs_parfirst : begin
574578
# n_θM, n_site_batch = size(ζMs)
575579
# get_concrete(ComponentArrayInterpreter(
576580
# P = (n_MC, n_θP), Ms = (n_MC, n_θM, n_site_batch)))
577581
# end
578582
#urandn = _create_randn(rng, CA.getdata(ζP), n_MC, n_θP + n_θMs)
579-
z = _create_randn(rng, CA.getdata(ζP), n_MC, n_θP + n_θMs)
580-
zP = @view z[:, 1:n_θP]
581-
zMs = @view z[:, (n_θP+1):end]
582-
sample_ζresid_norm(approx, zP, zMs, CA.getdata(ζP), CA.getdata(ζMs), args...;
583+
#z = _create_randn(rng, CA.getdata(ζP), n_MC, n_θP)
584+
zP = _create_randn(rng, CA.getdata(ζP), n_MC, n_θP)
585+
zMs = _create_randn(rng, CA.getdata(ζP), n_MC, n_θMs)
586+
sample_ζresid_norm(approx, zP, zMs, CA.getdata(ϕm), ϕq, args...;
583587
cor_ends,
584588
int_ϕq=get_concrete(int_ϕq)
585589
)
586590
end
587591

588-
function sample_ζresid_norm(approx::MeanHVIApproximationMat, zP::AbstractMatrix, zMs::AbstractMatrix,
589-
ζP::TP, ζMs::TM, ϕq::AbstractVector;
592+
function sample_ζresid_norm(approx::MeanHVIApproximationMat,
593+
zP::AbstractMatrix, zMs::AbstractMatrix,
594+
ϕm::TM, ϕq::AbstractVector{T};
590595
int_ϕq=get_concrete(ComponentArrayInterpreter(ϕq)),
591596
cor_ends
592-
) where {T,TP<:AbstractVector{T},TM<:AbstractMatrix{T}}
593-
ϕuncc = int_ϕq(CA.getdata(ϕq))
597+
) where {T,TM<:AbstractMatrix{T}}
598+
ζMs = ϕm
599+
ϕuncc = ϕqc = int_ϕq(CA.getdata(ϕq))
600+
ζP = ϕqc[Val(:μP)]
594601
n_θP, n_θMs, (n_θM, n_batch) = length(ζP), length(ζMs), size(ζMs)
595602
# do not create a UpperTriangular Matrix of an AbstractGÜUArray in transformU_cholesky1
596603
ρsP = isempty(ϕuncc[Val(:ρsP)]) ? similar(ϕuncc[Val(:ρsP)]) : ϕuncc[Val(:ρsP)] # required by zygote

‎src/elbo2.jl‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -16,13 +16,15 @@
1616
# and forward runs for many sites (760mus cpu vs 439ms gpu unit!)
1717

1818
function sample_ζresid_norm(app::MeanHVIApproximation,
19-
zP::AbstractMatrix, zMs::AbstractMatrix, ζP::TP, ζMs::TM,
20-
ϕq::AbstractVector;
19+
zP::AbstractMatrix, zMs::AbstractMatrix,
20+
ϕm::TM, ϕq::AbstractVector{T};
2121
int_ϕq=get_concrete(ComponentArrayInterpreter(ϕq)),
2222
cor_ends
23-
) where {T,TP<:AbstractVector{T},TM<:AbstractMatrix{T}}
23+
) where {T,TM<:AbstractMatrix{T}}
24+
ζMs = ϕm
2425
ϕunc = CA.getdata(ϕq)
25-
ϕuncc = int_ϕq(ϕunc)
26+
ϕuncc = ϕqc = int_ϕq(ϕunc)
27+
ζP = ϕqc[Val(:μP)]
2628
n_θP, n_θMs, (n_θM, n_batch), n_MC = length(ζP), length(ζMs), size(ζMs), size(zP,1)
2729
# do not create a UpperTriangular Matrix of an AbstractGÜUArray in transformU_cholesky1
2830
ρsP = isempty(ϕuncc[Val(:ρsP)]) ? similar(ϕuncc[Val(:ρsP)]) : ϕuncc[Val(:ρsP)] # required by zygote

‎src/elbo_dev.jl‎

Lines changed: 0 additions & 89 deletions
Original file line numberDiff line numberDiff line change
@@ -1,89 +0,0 @@
1-
2-
function sample_ζresid_norm(app::MeanHVIApproximationDev,
3-
zP::AbstractMatrix, zMs::AbstractMatrix, ζP::TP, ζMs::TM,
4-
ϕq::AbstractVector;
5-
int_ϕq=get_concrete(ComponentArrayInterpreter(ϕq)),
6-
cor_ends
7-
) where {T,TP<:AbstractVector{T},TM<:AbstractMatrix{T}}
8-
ϕunc = CA.getdata(ϕq)
9-
ϕuncc = int_ϕq(ϕunc)
10-
n_θP, n_θMs, (n_θM, n_batch), n_MC = length(ζP), length(ζMs), size(ζMs), size(zP,1)
11-
# do not create a UpperTriangular Matrix of an AbstractGÜUArray in transformU_cholesky1
12-
ρsP = isempty(ϕuncc[Val(:ρsP)]) ? similar(ϕuncc[Val(:ρsP)]) : ϕuncc[Val(:ρsP)] # required by zygote
13-
UPs, rangesP = transformU_blocks_cholesky1(ρsP, cor_ends.P)
14-
ρsM = isempty(ϕuncc[Val(:ρsM)]) ? similar(ϕuncc[Val(:ρsM)]) : ϕuncc[Val(:ρsM)] # required by zygote
15-
# cholesky factor of the correlation: diag(UM' * UM) .== 1
16-
# coefficients ρsM can be larger than 1, still yielding correlations <1 in UM' * UM
17-
UMs, rangesM = transformU_blocks_cholesky1(ρsM, cor_ends.M)
18-
cf = ϕuncc[Val(:coef_logσ2_ζMs)]
19-
logσ2_logMs = vec(cf[1, :] .+ cf[2, :] .* ζMs)
20-
logσ2_ζP = vec(CA.getdata(ϕuncc[Val(:logσ2_ζP)]))
21-
# CUDA cannot multiply BlockDiagonal * Diagonal, construct already those blocks
22-
σMs = reshape(exp.(logσ2_logMs ./ 2), n_θM, :)
23-
σP = exp.(logσ2_ζP ./ 2)
24-
# create random numbers from U diag(σ) z = U (σ .* z)
25-
# for each block separately
26-
#Ui, ri = first(zip(UPs, rangesP))
27-
ζP1 = ChainRulesCore.@ignore_derivatives similar(CA.getdata(ϕunc), size(zP'))
28-
#@benchmark typeof(zP .* ϕq[1:1]) # more allocations and more time
29-
local ζP_resids = if isempty(ζP)
30-
ζP1 # provide init of correct empty matrix type
31-
else
32-
mapreduce(vcat, UPs, rangesP) do Ui, ri
33-
Diagonal(σP[ri]) * Ui' * zP[:,ri]'
34-
end
35-
end::typeof(ζP1)
36-
#
37-
#Ui, ri = first(zip(UMs, rangesM))
38-
zMs_subjects = reshape(zMs, (n_MC, n_θM, n_batch))
39-
#σM, zM = first(zip(eachcol(σMs), eachslice(zMs_subjects; dims=3)))
40-
cat3 = (x,y) -> cat(x,y,dims=3)
41-
# map across subjects (n_batch)
42-
#ζMs_vec = map(eachcol(σMs), eachslice(zMs_subjects; dims=3)) do σM, zM
43-
ζMs1 = ChainRulesCore.@ignore_derivatives similar(ϕunc, size(zMs_subjects[:,:,1]'))
44-
fBlock = let UMs = UMs, rangesM = collect(rangesM) # without collect, type unstable
45-
function fBlock_inner(σM, zM)::typeof(ζMs1)
46-
# generator does not work with Zygote
47-
# vcat(...) is not faster than mapreduce
48-
# g = map( UMs, rangesM) do Ui, ri # n_θM, n_MC
49-
# #(zM[:,ri]' * Ui * diagm(σM[ri]))'
50-
# diagm(σM[ri]) * Ui' * zM[:, ri]'
51-
# end
52-
# vcat(g...)
53-
mapreduce(vcat, UMs, rangesM) do Ui, ri # n_θM, n_MC
54-
# (zM[:,ri]' * (Ui * Diagonal(σM[ri])))'
55-
# Diagonal is faster than diagm and works with GPU
56-
Diagonal(σM[ri]) * Ui' * zM[:, ri]' # transpose to not return adjoint
57-
end
58-
end
59-
end
60-
#ζMs1 = ChainRulesCore.@ignore_derivatives fBlock(σMs[:,1], zMs_subjects[:,:,1])
61-
#TMR = ChainRulesCore.@ignore_derivatives typeof(similar(ϕq, size(zMs)))
62-
#TB = Base.infer_return_type(fBlock, (typeof(σMs[:,1]), typeof(zMs_subjects[:,:,1])))
63-
#ζMs_vec = map(fBlock, eachcol(σMs[:,2:end]), eachslice(zMs_subjects[:,:,2:end]; dims=3), init = ζMs1)
64-
#ζMs_vec = map((σM, zM) -> fBlock(σM, zM)::TMR, eachcol(σMs), eachslice(zMs_subjects; dims=3))
65-
ζMs_vec = map((σM, zM) -> fBlock(σM, zM), eachcol(σMs), eachslice(zMs_subjects; dims=3))
66-
#zM = zMs_subjects[:,:,1]
67-
#ζMs_vec = [fBlock(σMs[:,i], zM) for i in axes(σMs, 2)]
68-
#ζMs_vec = [fBlock(σMs[:,1], zM)]
69-
#ζMs_vec = [fBlock(σMs[:,1], zM) for i in 1:n_batch]
70-
#ζMs_vec = [fBlock(σMs[:,1], zM) for i in axes(σMs, 2)]
71-
#ζMs_vec = [fBlock(σM, zM) for (σM, zM) in zip(eachcol(σMs), eachslice(zMs_subjects; dims=3))]
72-
# concatenate so that n_MC is last dimension
73-
local ζMs_parfirst_resids = stack(ζMs_vec; dims = 2 ) # n_θM, n_batch, n_MC
74-
#size(ζMs_parfirst_resids)
75-
76-
# std(ζMs_parfirst_resids[1,1,:])
77-
# std(ζMs_parfirst_resids[1,end,:])
78-
# σzMs_stacked = reshape(σzMs, (n_θM, n_batch * n_MC))
79-
# ζMs_resids_stacked = mapreduce(vcat, UMs, rangesM) do Ui, ri
80-
# #Ui * σzMs_stacked[ri, :]
81-
# Uσ = Ui * σMs[ri,:]
82-
# Uσ' * zMs[:,ri]'
83-
# diagm(σMs[ri,:]) * Ui' * zMs[:,ri]'
84-
# end
85-
# ζMs_parfirst_resids = reshape(ζMs_resids_stacked, n_θM, n_batch, n_MC)
86-
#
87-
diagUσ = vcat(σP, vec(σMs))
88-
ζP_resids, ζMs_parfirst_resids, diagUσ
89-
end

‎src/gf.jl‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -182,8 +182,8 @@ Provide a `transMs = StackedArray(transM, n_batch)`
182182
function gtrans(g, transMs, xMP, ϕg; cdev, is_testmode)
183183
# TODO remove after removing gf
184184
# predict the log of the parameters
185-
ζMst = g(xMP, ϕg; is_testmode)
186-
ζMs = ζMst'
185+
ϕg = g(xMP, ϕg; is_testmode)
186+
ζMs = ϕg'
187187
ζMs_cpu = cdev(ζMs)
188188
θMs = transMs(ζMs_cpu)
189189
if !all(isfinite.(θMs))

‎test/test_sample_zeta.jl‎

Lines changed: 33 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -38,23 +38,26 @@ function test_with_scenario(scenario)
3838
cor_ends = (P=(n_θP == 0 ? [0] : 1:n_θP), M=[n_θM])
3939
ρsP = zeros(FT, get_cor_count(cor_ends.P)) .+ FT(0.02)
4040
ρsM = zeros(FT, get_cor_count(cor_ends.M)) .+ FT(0.02)
41+
(; transP, transM) = get_hybridproblem_transforms(prob; scenario)
4142

4243
ϕq = CA.ComponentVector(;
4344
logσ2_ζP=fill(FT(-10.0), n_θP),
4445
coef_logσ2_ζMs=reduce(hcat, (FT[-10.0, 0.0] for _ in 1:n_θM)),
4546
ρsP,
46-
ρsM)
47+
ρsM,
48+
)
49+
50+
ϕq_true = CP.update_μP_by_θP(ϕq, θP_true, transP)
51+
ϕq = CA.ComponentVector(ϕq_true; μP = ϕq_true.μP .+ FT(0.01))
4752

4853
θ_true = θ = CA.ComponentVector(;
4954
P=θP_true,
5055
Ms=θMs_true)
51-
transPMs = elementwise(exp) # all parameters on LogNormal scale
52-
ζ_true = inverse(transPMs)(θ_true)
53-
ϕ_true = vcat(ζ_true, CA.ComponentVector(ϕq=ϕq))
54-
ϕ_cpu = vcat(ζ_true .+ FT(0.01), CA.ComponentVector(ϕq=ϕq))
56+
ϕ_true = CA.ComponentVector(Ms = θMs_true, ϕq=ϕq_true)
57+
ϕ_cpu = CA.ComponentVector(Ms = θMs_true .+ FT(0.01), ϕq=ϕq)
5558

5659
interpreters = (; pmu=ComponentArrayInterpreter(ϕ_true),
57-
ϕq=ComponentArrayInterpreter(ϕ_true.ϕq)
60+
ϕq=ComponentArrayInterpreter(ϕq)
5861
) #, M=int_θM, PMs=int_θPMs)
5962

6063
n_MC = 3
@@ -88,38 +91,37 @@ function test_with_scenario(scenario)
8891
# approx = MeanHVIApproximationMat()
8992
# approx = CP.MeanHVIApproximationDev()
9093
function test_sample_ζresid_norm(approx)
91-
ϕ = CA.getdata(ϕ_cpu)
92-
ϕc = interpreters.pmu(ϕ)
94+
ϕc = copy(ϕ_cpu)
9395
ϕc.ϕq.coef_logσ2_ζMs[1,:] .= (log ∘ abs2).((0.1, 100.0))
9496
ϕc.ϕq.ρsM .= 0.0
9597
int_ϕq = get_concrete(ComponentArrayInterpreter(ϕc.ϕq))
9698
n_MC_pred = 300 # larger n_MC to test σ2
9799
n_site_batch = size(ϕc.Ms,2)
98100
#rng = StableRNG(111)
99101
# @inferred gives any, while Cthulhu inferres concrete type
100-
# ζP_resids, ζMs_parfirst_resids, σ = @inferred CP.sample_ζresid_norm(approx, rng, ϕc.P, ϕc.Ms, ϕc.ϕq;
102+
# ζP_resids, ζMs_parfirst_resids, σ = @inferred CP.sample_ζresid_norm(approx, rng, ϕc.Ms, ϕc.ϕq;
101103
# n_MC=n_MC_pred, cor_ends, int_ϕq)
102-
# @inferred first(CP.sample_ζresid_norm(approx, rng, ϕc.P, ϕc.Ms, ϕc.ϕq;
104+
# @inferred first(CP.sample_ζresid_norm(approx, rng, ϕc.Ms, ϕc.ϕq;
103105
# n_MC=n_MC_pred, cor_ends, int_ϕq))
104106
# ζP_resids, ζMs_parfirst_resids, σ = CP.sample_ζresid_norm(approx, rng, ϕc.P, ϕc.Ms, ϕc.ϕq;
105107
# n_MC=n_MC_pred, cor_ends, int_ϕq)
106-
ζP_resids, ζMs_parfirst_resids, σ = @inferred CP.sample_ζresid_norm(approx, rng, ϕc.P, ϕc.Ms, ϕc.ϕq;
108+
ζP_resids, ζMs_parfirst_resids, σ = @inferred CP.sample_ζresid_norm(approx, rng, ϕc.Ms, ϕc.ϕq;
107109
n_MC=n_MC_pred, cor_ends, int_ϕq)
108-
#@code_warntype CP.sample_ζresid_norm(approx, rng, ϕc.P, ϕc.Ms, ϕc.ϕq; n_MC=n_MC_pred, cor_ends, int_ϕq)
110+
#@code_warntype CP.sample_ζresid_norm(approx, rng, ϕc.Ms, ϕc.ϕq; n_MC=n_MC_pred, cor_ends, int_ϕq)
109111
#@usingany Cthulhu
110-
#@descend_code_warntype CP.sample_ζresid_norm(approx, rng, ϕc.P, ϕc.Ms, ϕc.ϕq; n_MC=n_MC_pred, cor_ends, int_ϕq)
112+
#@descend_code_warntype CP.sample_ζresid_norm(approx, rng, ϕc.Ms, ϕc.ϕq; n_MC=n_MC_pred, cor_ends, int_ϕq)
111113
#@test size(ζ_resid) == (length(ϕc.P) + n_site * n_θM, n_MC)
112114
n_θM = size(ϕc.Ms,1)
113115
@test size(ζP_resids) == (n_θP, n_MC_pred)
114116
@test size(ζMs_parfirst_resids) == (n_θM, n_site_batch, n_MC_pred)
115117
gr =
116118
Zygote.gradient(ϕc -> begin
117119
ζP_resids, ζMs_parfirst_resids, σ = CP.sample_ζresid_norm(
118-
approx, rng, ϕc.P, ϕc.Ms, ϕc.ϕq;
120+
approx, rng, ϕc.Ms, ϕc.ϕq;
119121
n_MC, cor_ends, int_ϕq)
120122
sum(ζP_resids) + sum(ζMs_parfirst_resids)
121123
end, ϕc)[1]
122-
@test length(gr) == length(ϕ)
124+
@test length(gr) == length(ϕc)
123125
#
124126
n_θM, n_site_batch = size(ϕc.Ms)
125127
# intm_PMs = ComponentArrayInterpreter(
@@ -139,32 +141,35 @@ function test_with_scenario(scenario)
139141
#ζP, ζMs, ϕq = ϕc.P, ϕc.Ms, ϕc.ϕq
140142
#urandn = CUDA.randn(length(ϕc.P) + length(ϕc.Ms), n_MC) |> gpu
141143
#include(joinpath(@__DIR__, "uncNN", "elbo.jl")) # callback_loss
142-
#ζ_resid, σ = sample_ζresid_norm(urandn, ϕc.P, ϕc.Ms, ϕc.ϕq; n_MC)
143-
#Zygote.gradient(ϕc -> sum(sample_ζresid_norm(urandn, ϕc.P, ϕc.Ms, ϕc.ϕq; n_MC)[1]), ϕc)[1];
144+
#ζ_resid, σ = sample_ζresid_norm(urandn, ϕc.Ms, ϕc.ϕq; n_MC)
145+
#Zygote.gradient(ϕc -> sum(sample_ζresid_norm(urandn, ϕc.Ms, ϕc.ϕq; n_MC)[1]), ϕc)[1];
144146
# @inferred first(CP.sample_ζresid_norm(
145-
# approx, rng, CA.getdata(ϕcd.P), CA.getdata(ϕcd.Ms), CA.getdata(ϕcd.ϕq);
147+
# approx, rng, CA.getdata(ϕcd.Ms), CA.getdata(ϕcd.ϕq);
146148
# n_MC = n_MC_pred, cor_ends, int_ϕq))
147149
# ζP_resids, ζMs_parfirst_resids, σ = CP.sample_ζresid_norm(
148-
# approx, rng, CA.getdata(ϕcd.P), CA.getdata(ϕcd.Ms), CA.getdata(ϕcd.ϕq);
150+
# approx, rng, CA.getdata(ϕcd.Ms), CA.getdata(ϕcd.ϕq);
149151
# n_MC = n_MC_pred, cor_ends, int_ϕq)
150152
ζP_resids, ζMs_parfirst_resids, σ = @inferred CP.sample_ζresid_norm(
151-
approx, rng, CA.getdata(ϕcd.P), CA.getdata(ϕcd.Ms), CA.getdata(ϕcd.ϕq);
153+
approx, rng, CA.getdata(ϕcd.Ms), CA.getdata(ϕcd.ϕq);
152154
n_MC = n_MC_pred, cor_ends, int_ϕq)
153-
#@descend_code_warntype CP.sample_ζresid_norm(rng, CA.getdata(ϕcd.P), CA.getdata(ϕcd.Ms), CA.getdata(ϕcd.ϕq); n_MC = n_MC_pred, cor_ends, int_ϕq)
155+
#@descend_code_warntype CP.sample_ζresid_norm(rng, CA.getdata(ϕcd.Ms), CA.getdata(ϕcd.ϕq); n_MC = n_MC_pred, cor_ends, int_ϕq)
154156
@test ζP_resids isa GPUArraysCore.AbstractGPUArray
155157
@test ζMs_parfirst_resids isa GPUArraysCore.AbstractGPUArray
156158
@test size(ζP_resids) == (n_θP, n_MC_pred)
157159
@test size(ζMs_parfirst_resids) == (n_θM, n_site_batch, n_MC_pred)
158160
# Zygote gradient for many sites, use fewer sites here
159161
n_site_few = 20
160-
ϕcd_few = CA.ComponentVector(; P = ϕcd.P, Ms = ϕcd.Ms[:,1:n_site_few], ϕq = ϕcd.ϕq);
162+
# replacing Ms in ComponentVector by different length does not work on GPU
163+
# using workaround by NamedTuples, results in ScalarIndexing, need explicit
164+
# ϕcd_few = CA.ComponentVector(ϕcd; Ms = ϕcd.Ms[:,1:n_site_few]);
165+
ϕcd_few = CA.ComponentVector(Ms = ϕcd.Ms[:,1:n_site_few], ϕq = ϕcd.ϕq);
161166
#@usingany BenchmarkTools
162167
gr =
163168
#@profview Zygote.gradient(ϕc -> begin # type stable, most time spent in mapreduce
164169
#@benchmark Zygote.gradient(ϕc -> begin # many small allocs
165170
Zygote.gradient(ϕc -> begin
166171
ζP_resids, ζMs_parfirst_resids, σ = CP.sample_ζresid_norm(
167-
approx, rng, CA.getdata(ϕc.P), CA.getdata(ϕc.Ms), CA.getdata(ϕc.ϕq);
172+
approx, rng, CA.getdata(ϕc.Ms), CA.getdata(ϕc.ϕq);
168173
n_MC, cor_ends, int_ϕq)
169174
sum(ζP_resids) + sum(ζMs_parfirst_resids)
170175
end, ϕcd_few)[1]; # semicolon required
@@ -175,26 +180,26 @@ function test_with_scenario(scenario)
175180
() -> begin
176181
CP.sample_ζresid_norm(
177182
#@benchmark CP.sample_ζresid_norm(
178-
approx, rng, ϕc.P, ϕc.Ms, ϕc.ϕq;
183+
approx, rng, ϕc.Ms, ϕc.ϕq;
179184
n_MC, cor_ends, int_ϕq)
180185
#
181186
CP.sample_ζresid_norm(
182187
#@benchmark CP.sample_ζresid_norm(
183-
approx, rng, ϕcd.P, ϕcd.Ms, ϕcd.ϕq;
188+
approx, rng, ϕcd.Ms, ϕcd.ϕq;
184189
n_MC, cor_ends, int_ϕq)
185190
#
186-
ϕc_few = CA.ComponentVector(; P = ϕc.P, Ms = ϕc.Ms[:,1:n_site_few], ϕq = ϕc.ϕq);
191+
ϕc_few = CA.ComponentVector(ϕc; Ms = ϕc.Ms[:,1:n_site_few]);
187192
Zygote.gradient(ϕc -> begin
188193
#@benchmark Zygote.gradient(ϕc -> begin # many small allocs
189194
ζP_resids, ζMs_parfirst_resids, σ = CP.sample_ζresid_norm(
190-
approx, rng, ϕc.P, ϕc.Ms, ϕc.ϕq;
195+
approx, rng, ϕc.Ms, ϕc.ϕq;
191196
n_MC, cor_ends, int_ϕq)
192197
sum(ζP_resids) + sum(ζMs_parfirst_resids)
193198
end, ϕc_few)[1]
194199
Zygote.gradient(ϕc -> begin # many small allocs
195200
#@benchmark Zygote.gradient(ϕc -> begin # many small allocs
196201
ζP_resids, ζMs_parfirst_resids, σ = CP.sample_ζresid_norm(
197-
approx, rng, CA.getdata(ϕc.P), CA.getdata(ϕc.Ms), CA.getdata(ϕc.ϕq);
202+
approx, rng, CA.getdata(ϕc.Ms), CA.getdata(ϕc.ϕq);
198203
n_MC, cor_ends, int_ϕq)
199204
sum(ζP_resids) + sum(ζMs_parfirst_resids)
200205
end, ϕcd_few)[1]

0 commit comments

Comments
 (0)