move PBMApplicator to gdev_P int solve and predict_hvi

replace stack in apply_process_model by mapreduce

Author: bgctw

src/HybridSolver.jl

@@ -134,6 +134,12 @@ function CommonSolve.solve(prob::AbstractHybridProblem, solver::HybridPosteriorS
         train_loader_dev = train_loader
     end
     f = get_hybridproblem_PBmodel(prob; scenario, use_all_sites=false)
+    if gdevs.gdev_P isa MLDataDevices.AbstractGPUDevice
+        f_dev = fmap(gdevs.gdev_P, f)
+    else
+        f_dev = f
+    end
+
     py = get_hybridproblem_neg_logden_obs(prob; scenario)
 
     priors_θP_mean, priors_θMs_mean = construct_priors_θ_mean(
@@ -142,7 +148,7 @@ function CommonSolve.solve(prob::AbstractHybridProblem, solver::HybridPosteriorS
     y_global_o = Float32[] # TODO
 
     loss_elbo = get_loss_elbo(
-        g_dev, transP, transMs, f, py, y_global_o;
+        g_dev, transP, transMs, f_dev, py, y_global_o;
         solver.n_MC, solver.n_MC_cap, cor_ends, priors_θP_mean, priors_θMs_mean, 
         cdev=infer_cdev(gdevs), pbm_covars, θP, int_unc, int_μP_ϕg_unc)
     # test loss function once

src/HybridVariationalInference.jl

@@ -84,7 +84,7 @@ include("gencovar.jl")
 export callback_loss
 include("util_opt.jl")
 
-export cpu_ca
+export cpu_ca, apply_preserve_axes
 include("util_ca.jl")
 
 export neg_logden_indep_normal, entropy_MvNormal

src/elbo.jl

@@ -172,7 +172,12 @@ function predict_hvi(rng, prob::AbstractHybridProblem; scenario=Val(()),
     is_predict_batch = (n_site_pred == n_batch)
     @assert size(xP, 2) == n_site_pred
     f = get_hybridproblem_PBmodel(prob; scenario, use_all_sites=!is_predict_batch)
-    y = apply_process_model(θsP, θsMs, f, xP)
+    if gdevs.gdev_P isa MLDataDevices.AbstractGPUDevice
+        f_dev = fmap(gdevs.gdev_P, f)
+    else
+        f_dev = f
+    end
+    y = apply_process_model(θsP, θsMs, f_dev, xP)
     (; y, θsP, θsMs, entropy_ζ)
 end
 
@@ -317,10 +322,20 @@ Call a PBM applicator for a sample of parameters of each site, and stack results
 Results are of shape `(n_obs x n_site_pred x n_MC)`.
 """
 function apply_process_model(θsP::AbstractMatrix, θsMs::AbstractArray{ET,3}, f, xP) where ET
-    y_pred = stack(map(eachcol(θsP), eachslice(θsMs, dims=3)) do θP, θMs
+    # stack does not work on GPU
+    # y_pred = stack(map(eachcol(θsP), eachslice(θsMs, dims=3)) do θP, θMs
+    #     y_global, y_pred_i = f(θP, θMs, xP)
+    #     y_pred_i
+    # end)
+    #Main.@infiltrate_main
+    # for type stability, apply f at first sample before mapreduce
+    P1, Pit = Iterators.peel(eachcol(θsP));
+    Ms1, Msit = Iterators.peel(eachslice(θsMs, dims=3));
+    y1 = f(P1, Ms1, xP)[2]
+    y_pred = mapreduce((a,b) -> cat(a,b;dims=3), Pit, Msit; init=y1) do θP, θMs
         y_global, y_pred_i = f(θP, θMs, xP)
         y_pred_i
-    end)
+    end
 end
 
 """

src/util_ca.jl

@@ -3,13 +3,17 @@
 
 Move ComponentArray form gpu to cpu.    
 """
-#function cpu_ca end
-# define in FluxExt
 function cpu_ca(ca::CA.ComponentArray)
     CA.ComponentArray(cpu_device()(CA.getdata(ca)), CA.getaxes(ca))
 end
 
+"""
+    apply_preserve_axes(f, ca::ComponentArray)
 
+Apply callable `f(x)` to the data inside `ca`, assume that the result has
+the same shape, and return a new `ComponentArray` with the same axes
+as in `ca`.
+"""
 function apply_preserve_axes(f, ca::CA.ComponentArray)
     CA.ComponentArray(f(CA.getdata(ca)), CA.getaxes(ca))
 end