HybirdProblem constructors with SimpleChain, Flux.Chain and Lux.Chain

Author: bgctw

ext/HybridVariationalInferenceFluxExt.jl

@@ -2,6 +2,7 @@ module HybridVariationalInferenceFluxExt
 
 using HybridVariationalInference, Flux
 using HybridVariationalInference: HybridVariationalInference as HVI
+using ComponentArrays: ComponentArrays as CA
 
 struct FluxApplicator{RT} <: AbstractModelApplicator
     rebuild::RT
@@ -25,6 +26,14 @@ function __init__()
     HVI.set_default_GPUHandler(FluxGPUDataHandler())
 end
 
+function HVI.HybridProblem(θP::CA.ComponentVector, θM::CA.ComponentVector, g_chain::Flux.Chain, 
+    args...; kwargs...)
+    # constructor with Flux.Chain
+    ϕ, _ = destructure(g_chain)
+    g = construct_FluxApplicator(g_chain), ϕ
+    HybridProblem(θP, θM, g, ϕg, args...; kwargs...)
+end
+
 function HVI.get_hybridcase_MLapplicator(case::HVI.DoubleMM.DoubleMMCase, ::Val{:Flux};
         scenario::NTuple = ())
     (; n_covar, n_θM) = get_hybridcase_sizes(case; scenario)
@@ -43,4 +52,6 @@ function HVI.get_hybridcase_MLapplicator(case::HVI.DoubleMM.DoubleMMCase, ::Val{
     construct_FluxApplicator(g_chain), ϕ
 end
 
+
+
 end # module

ext/HybridVariationalInferenceLuxExt.jl

@@ -25,4 +25,13 @@ function HVI.apply_model(app::LuxApplicator, x, ϕ)
     app.stateful_layer(x, ϕc)
 end
 
+function HVI.HybridProblem(θP::CA.ComponentVector, θM::CA.ComponentVector, g_chain::Chain, 
+    args...; kwargs...)
+    # constructor with SimpleChain
+    g = construct_LuxApplicator(g_chain)
+    FT = eltype(θM)
+    ϕg = randn(FT, length(g.int_ϕ))
+    HybridProblem(θP, θM, g, ϕg, args...; kwargs...)
+end
+
 end # module

ext/HybridVariationalInferenceSimpleChainsExt.jl

@@ -3,6 +3,9 @@ module HybridVariationalInferenceSimpleChainsExt
 using HybridVariationalInference, SimpleChains
 using HybridVariationalInference: HybridVariationalInference as HVI
 using StatsFuns: logistic
+using ComponentArrays: ComponentArrays as CA
+
+
 
 struct SimpleChainsApplicator{MT} <: AbstractModelApplicator
     m::MT
@@ -12,6 +15,14 @@ HVI.construct_SimpleChainsApplicator(m::SimpleChain) = SimpleChainsApplicator(m)
 
 HVI.apply_model(app::SimpleChainsApplicator, x, ϕ) = app.m(x, ϕ)
 
+function HVI.HybridProblem(θP::CA.ComponentVector, θM::CA.ComponentVector, g_chain::SimpleChain, 
+    args...; kwargs...)
+    # constructor with SimpleChain
+    g = construct_SimpleChainsApplicator(g_chain)
+    ϕg = SimpleChains.init_params(g_chain, eltype(θM))
+    HybridProblem(θP, θM, g, ϕg, args...; kwargs...)
+end
+
 function HVI.get_hybridcase_MLapplicator(case::HVI.DoubleMM.DoubleMMCase, ::Val{:SimpleChains};
         scenario::NTuple=())
     (;n_covar, n_θM) = get_hybridcase_sizes(case; scenario)

src/HybridProblem.jl

@@ -12,10 +12,12 @@ struct HybridProblem <: AbstractHybridCase
     # inner constructor to constrain the types
     function HybridProblem(
         θP::CA.ComponentVector, θM::CA.ComponentVector, 
+        g::AbstractModelApplicator, ϕg, 
+        f::Function, 
         transM::Union{Function, Bijectors.Transform}, 
         transP::Union{Function, Bijectors.Transform}, 
         n_covar::Integer, n_batch::Integer, 
-        f::Function, g::AbstractModelApplicator, ϕg, train_loader::DataLoader)
+        train_loader::DataLoader)
         new(θP, θM, transM, transP, n_covar, n_batch, f, g, ϕg, train_loader)
     end
 end

test/test_HybridProblem.jl

@@ -44,20 +44,21 @@ construct_problem = () -> begin
         # dense layer without bias that maps to n outputs and `identity` activation
         TurboDense{false}(identity, n_out)
     )
-    g = construct_SimpleChainsApplicator(g_chain)
-    ϕg = SimpleChains.init_params(g_chain, eltype(θM))
+    # g = construct_SimpleChainsApplicator(g_chain)
+    # ϕg = SimpleChains.init_params(g_chain, eltype(θM))
     #
     rng = StableRNG(111)
     (; xM, n_site, θP_true, θMs_true, xP, y_global_true, y_true, y_global_o, y_o
-    ) = gen_hybridcase_synthetic(DoubleMM.DoubleMMCase(), rng;);
+) = gen_hybridcase_synthetic(DoubleMM.DoubleMMCase(), rng;)
     train_loader = MLUtils.DataLoader((xM, xP, y_o), batchsize = n_batch)
-    HybridProblem(θP, θM, transM, transP, n_covar, n_batch, f_doubleMM_with_global, 
-        g, ϕg, train_loader)
+    # HybridProblem(θP, θM, transM, transP, n_covar, n_batch, f_doubleMM_with_global, 
+    #     g, ϕg, train_loader)
+    HybridProblem(θP, θM, g_chain, f_doubleMM_with_global, 
+        transM, transP, n_covar, n_batch, train_loader)
 end
 prob = construct_problem();
 scenario = (:default,)
 
-
 #(; n_covar, n_batch, n_θM, n_θP) = get_hybridcase_sizes(prob; scenario)
 
 @testset "loss_gf" begin