implement and test a type-stable rebuild FluxApplicator

test show no significant speedup

Author: bgctw

ext/HybridVariationalInferenceFluxExt.jl

@@ -10,6 +10,13 @@ struct FluxApplicator{RT} <: AbstractModelApplicator
     rebuild::RT
 end
 
+struct PartricFluxApplicator{RT, MT, YT} <: AbstractModelApplicator
+    rebuild::RT
+end
+
+const FluxApplicatorU{RT} = Union{FluxApplicator{RT},PartricFluxApplicator{RT}} where RT
+
+
 function HVI.construct_ChainsApplicator(rng::AbstractRNG, m::Chain, float_type::DataType)
     # TODO: care fore rng and float_type
     ϕ, rebuild = Flux.destructure(m)
@@ -23,7 +30,18 @@ function HVI.apply_model(app::FluxApplicator, x::T, ϕ) where T
     res = m(x)::T
     res
 end
-function _apply_model(m,x) # function barrier so that m is inferred
+
+
+function HVI.construct_partric(app::FluxApplicator{RT}, x, ϕ) where RT
+    m = app.rebuild(ϕ)
+    y = m(x)
+    PartricFluxApplicator{RT, typeof(m), typeof(y)}(app.rebuild)
+end
+
+function HVI.apply_model(app::PartricFluxApplicator{RT, MT, YT}, x, ϕ) where {RT, MT, YT}
+    m = app.rebuild(ϕ)::MT
+    res = m(x)::YT
+    res
 end
 
 # struct FluxGPUDataHandler <: AbstractGPUDataHandler end
@@ -70,4 +88,5 @@ end
 
 
 
+
 end # module

src/HybridVariationalInference.jl

@@ -28,6 +28,7 @@ include("bijectors_utils.jl")
 export AbstractComponentArrayInterpreter, ComponentArrayInterpreter,
        StaticComponentArrayInterpreter
 export flatten1, get_concrete, get_positions, stack_ca_int, compose_interpreters
+export construct_partric
 include("ComponentArrayInterpreter.jl")
 
 export AbstractModelApplicator, construct_ChainsApplicator

src/ModelApplicator.jl

@@ -34,6 +34,16 @@ function apply_model(app::NullModelApplicator, x, ϕ)
     return x
 end
 
+"""
+Construct a parametric type-stable model applicator, given
+covariates, `x`, and parameters, `ϕ`.
+
+The default retuns the current model applicator.
+"""
+function construct_partric(app::AbstractModelApplicator, x, ϕ) 
+    app
+end
+
 
 """
     construct_ChainsApplicator([rng::AbstractRNG,] chain, float_type)

test/test_Flux.jl

@@ -56,6 +56,14 @@ using Flux
     # @descend_code_warntype g(x, ϕ)
     #@test ϕ isa GPUArraysCore.AbstractGPUArray
     @test size(y) == (n_out, n_site)
+    gp = construct_partric(g, x, ϕ)
+    y2 = @inferred gp(x, ϕ)
+    @test y2 == y
+    () -> begin
+        # @usingany BenchmarkTools
+        #@benchmark g(x,ϕ)
+        #@benchmark gp(x,ϕ) # no difference type-inferred
+    end
 end;
 
 @testset "cpu_ca" begin