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

Author: bgctw

ext/HybridVariationalInferenceFluxExt.jl

@@ -9,8 +9,8 @@ struct FluxApplicator{RT} <: AbstractModelApplicator
 end
 
 function HVI.construct_FluxApplicator(m::Chain)
-    _, rebuild = destructure(m)
-    FluxApplicator(rebuild)
+    ϕ, rebuild = destructure(m)
+    FluxApplicator(rebuild), ϕ
 end
 
 function HVI.apply_model(app::FluxApplicator, x, ϕ)
@@ -29,8 +29,7 @@ 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), ϕ
+    g, ϕg = construct_FluxApplicator(g_chain)
     HybridProblem(θP, θM, g, ϕg, args...; kwargs...)
 end
 
@@ -48,8 +47,7 @@ function HVI.get_hybridcase_MLapplicator(case::HVI.DoubleMM.DoubleMMCase, ::Val{
         # dense layer without bias that maps to n outputs and `identity` activation
         Flux.Dense(n_covar * 4 => n_out, identity, bias = false)
     )
-    ϕ, _ = destructure(g_chain)
-    construct_FluxApplicator(g_chain), ϕ
+    construct_FluxApplicator(g_chain)
 end
 
 

ext/HybridVariationalInferenceLuxExt.jl

@@ -10,14 +10,14 @@ struct LuxApplicator{MT, IT} <: AbstractModelApplicator
     int_ϕ::IT
 end
 
-function HVI.construct_LuxApplicator(m::Chain; device = gpu_device()) 
+function HVI.construct_LuxApplicator(m::Chain, float_type=Float32; device = gpu_device()) 
     ps, st = Lux.setup(Random.default_rng(), m)
-    ps_ca = CA.ComponentArray(ps) 
+    ps_ca = float_type.(CA.ComponentArray(ps)) 
     st = st |> device
     stateful_layer = StatefulLuxLayer{true}(m, nothing, st)
     #stateful_layer(x_o_gpu[:, 1:n_site_batch], ps_ca)
     int_ϕ = get_concrete(ComponentArrayInterpreter(ps_ca))
-    LuxApplicator(stateful_layer, int_ϕ)
+    LuxApplicator(stateful_layer, int_ϕ), ps_ca
 end
 
 function HVI.apply_model(app::LuxApplicator, x, ϕ) 
@@ -26,11 +26,9 @@ function HVI.apply_model(app::LuxApplicator, x, ϕ)
 end
 
 function HVI.HybridProblem(θP::CA.ComponentVector, θM::CA.ComponentVector, g_chain::Chain, 
-    args...; kwargs...)
+    args...; device = gpu_device(), kwargs...)
     # constructor with SimpleChain
-    g = construct_LuxApplicator(g_chain)
-    FT = eltype(θM)
-    ϕg = randn(FT, length(g.int_ϕ))
+    g, ϕg = construct_LuxApplicator(g_chain, eltype(θM); device)
     HybridProblem(θP, θM, g, ϕg, args...; kwargs...)
 end
 

ext/HybridVariationalInferenceSimpleChainsExt.jl

@@ -11,15 +11,17 @@ struct SimpleChainsApplicator{MT} <: AbstractModelApplicator
     m::MT
 end
 
-HVI.construct_SimpleChainsApplicator(m::SimpleChain) = SimpleChainsApplicator(m)
+function HVI.construct_SimpleChainsApplicator(m::SimpleChain, FloatType=Float32) 
+    ϕ = SimpleChains.init_params(m, FloatType);
+    SimpleChainsApplicator(m), ϕ
+end
 
 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))
+    g, ϕg = construct_SimpleChainsApplicator(g_chain)
     HybridProblem(θP, θM, g, ϕg, args...; kwargs...)
 end
 
@@ -50,8 +52,7 @@ function HVI.get_hybridcase_MLapplicator(case::HVI.DoubleMM.DoubleMMCase, ::Val{
             TurboDense{false}(identity, n_out)
         )
     end
-    ϕ = SimpleChains.init_params(g_chain, FloatType);
-    SimpleChainsApplicator(g_chain), ϕ
+    construct_SimpleChainsApplicator(g_chain, FloatType)
 end
 
 end # module

src/ModelApplicator.jl

@@ -1,3 +1,21 @@
+"""
+    AbstractModelApplicator(x, ϕ)
+
+Abstraction of applying a machine learning model at covariate matrix, `x`,
+using parameters, `ϕ`. It returns a matrix of predictions with the same
+number of rows as in `x`.    
+
+Constructors for specifics are defined in extension packages.
+Each constructor takes a special type of machine learning model and returns 
+a tuple with two components:
+- The applicator 
+- a sample parameter vector (type  depends on the used ML-framework)
+
+Implemented are
+- `construct_SimpleChainsApplicator`
+- `construct_FluxApplicator`
+- `construct_LuxApplicator`
+"""
 abstract type AbstractModelApplicator end
 
 function apply_model end

test/test_Flux.jl

@@ -35,16 +35,19 @@ end;
         Dense(n_covar * 4 => n_covar * 4, tanh),
         Dense(n_covar * 4 => n_out, identity, bias=false),
     )
-    g = construct_FluxApplicator(g_chain)
+    g, ϕg = construct_FluxApplicator(g_chain |> f64)
+    @test eltype(ϕg) == Float64
+    g, ϕg = construct_FluxApplicator(g_chain)
+    @test eltype(ϕg) == Float32
     n_site = 3
     x = rand(Float32, n_covar, n_site)
-    ϕ, _rebuild = destructure(g_chain)
-    y = g(x, ϕ)
+    #ϕ, _rebuild = destructure(g_chain)
+    y = g(x, ϕg)
     @test size(y) == (n_out, n_site)
     #
     n_site = 3
     x = rand(Float32, n_covar, n_site) |> gpu
-    ϕ = ϕ |> gpu
+    ϕ = ϕg |> gpu
     y = g(x, ϕ)
     #@test ϕ isa GPUArraysCore.AbstractGPUArray
     @test size(y) == (n_out, n_site)

test/test_HybridProblem.jl

@@ -44,8 +44,7 @@ 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, ϕg = construct_SimpleChainsApplicator(g_chain)
     #
     rng = StableRNG(111)
     (; xM, n_site, θP_true, θMs_true, xP, y_global_true, y_true, y_global_o, y_o

test/test_Lux.jl

@@ -1,8 +1,8 @@
 using HybridVariationalInference
 using Test
+using CUDA, GPUArraysCore
 using Lux
 using StatsFuns: logistic
-using CUDA, GPUArraysCore
 
 
 @testset "LuxModelApplicator" begin
@@ -13,18 +13,20 @@ using CUDA, GPUArraysCore
         Dense(n_covar * 4 => n_covar * 4, tanh),
         Dense(n_covar * 4 => n_out, logistic, use_bias=false),
     );
-    g = construct_LuxApplicator(g_chain; device = cpu_device());
+    g, ϕ = construct_LuxApplicator(g_chain, Float64; device = cpu_device());
+    @test eltype(ϕ) == Float64
+    g, ϕ = construct_LuxApplicator(g_chain; device = cpu_device());
+    @test eltype(ϕ) == Float32
     n_site = 3
     x = rand(Float32, n_covar, n_site)
-    ϕ = randn(Float32, Lux.parameterlength(g_chain))
+    #ϕ = randn(Float32, Lux.parameterlength(g_chain))
     y = g(x, ϕ)
     @test size(y) == (n_out, n_site)
     #
-    g = construct_LuxApplicator(g_chain; device = gpu_device());
-    n_site = 3
     x = rand(Float32, n_covar, n_site) |> gpu_device()
-    ϕ = randn(Float32, Lux.parameterlength(g_chain)) |> gpu_device()
-    y = g(x, ϕ)
+    ϕ_gpu = ϕ |> gpu_device()
+    #ϕ = randn(Float32, Lux.parameterlength(g_chain)) |> gpu_device()
+    y = g(x, ϕ_gpu)
     #@test ϕ isa GPUArraysCore.AbstractGPUArray
     @test size(y) == (n_out, n_site)
 end;

test/test_SimpleChains.jl

@@ -12,10 +12,10 @@ using StatsFuns: logistic
         TurboDense{true}(tanh, n_covar * 4),
         TurboDense{false}(logistic, n_out)
     )
-    g = construct_SimpleChainsApplicator(g_chain)
+    g, ϕg = construct_SimpleChainsApplicator(g_chain)
     n_site = 3
     x = rand(n_covar, n_site)
-    ϕ = SimpleChains.init_params(g_chain);
-    y = g(x, ϕ)
+    #ϕg = SimpleChains.init_params(g_chain);
+    y = g(x, ϕg)
     @test size(y) == (n_out, n_site)
 end;

test/test_cholesky_structure.jl

@@ -247,8 +247,8 @@ end
     #@test Upred ≈ CU
     SUpred = Upred * Dσ
     #hcat(SUpred, SU)  
-    @test SUpred≈SU atol=2e-1
+    @test SUpred≈SU atol=6e-1
     S_pred = Dσ' * Upred' * Upred * Dσ
-    @test S_pred≈S atol=2e-1
+    @test S_pred≈S atol=6e-1
 end