replace transformations by Exp

because exlementwise(exp) failed on AD on GPU
test on GPU

Author: bgctw

src/DoubleMM/f_doubleMM.jl

@@ -6,10 +6,10 @@ const θall = vcat(θP, θM)
 
 const θP_nor0 = θP[(:K2,)]
 
-const transP = elementwise(exp)
-const transM = elementwise(exp)
+# const transP = elementwise(exp)
+# const transM = elementwise(exp)
 
-const transMS = Stacked(elementwise(identity), elementwise(exp))
+# const transMS = Stacked(elementwise(identity), elementwise(exp))
 
 const int_θdoubleMM = ComponentArrayInterpreter(flatten1(CA.ComponentVector(; θP, θM)))
 
@@ -114,15 +114,17 @@ end
 
 
 function HVI.get_hybridproblem_transforms(prob::DoubleMMCase; scenario::NTuple = ())
+    _θP, _θM = get_hybridproblem_par_templates(prob; scenario)
     if (:stackedMS ∈ scenario)
-        return ((; transP, transM = transMS))
+        return (; transP = Stacked((HVI.Exp(),),(1:length(_θP),)),
+        transM = Stacked((identity,HVI.Exp(),),(1:1, 2:length(_θM),))) 
     elseif (:transIdent ∈ scenario)
         # identity transformations, should AD on GPU
-        _θP, _θM = get_hybridproblem_par_templates(prob; scenario)
         return (; transP = Stacked((identity,),(1:length(_θP),)),
             transM = Stacked((identity,),(1:length(_θM),)))
     end
-    (; transP, transM)
+    (; transP = Stacked((HVI.Exp(),),(1:length(_θP),)),
+    transM = Stacked((HVI.Exp(),),(1:length(_θM),)))    
 end
 
 # function HVI.get_hybridproblem_sizes(::DoubleMMCase; scenario = ())

src/HybridVariationalInference.jl

@@ -21,7 +21,7 @@ using StaticArrays: StaticArrays as SA
 using Functors
 
 #export Exp
-include("bijectors_exp.jl") 
+include("bijectors_utils.jl") 
 
 export ComponentArrayInterpreter, flatten1, get_concrete, get_positions
 include("ComponentArrayInterpreter.jl")

src/bijectors_exp.jl

@@ -0,0 +1,21 @@
+struct Exp <: Bijector
+end
+
+#Functors.@functor Exp
+Bijectors.transform(b::Exp, x) = exp.(x) # note the broadcast
+Bijectors.transform(ib::Inverse{<:Exp}, y) = log.(y)
+
+# `logabsdetjac`
+Bijectors.logabsdetjac(b::Exp, x) = sum(x)
+
+`with_logabsdet_jacobian`
+function Bijectors.with_logabsdet_jacobian(b::Exp, x)
+    return exp.(x), sum(x)
+end
+# function Bijectors.with_logabsdet_jacobian(ib::Inverse{<:Exp}, y)
+#     x = transform(ib, y)
+#     return x, -logabsdetjac(inverse(ib), x)
+# end
+
+
+Bijectors.is_monotonically_increasing(::Exp) = true

src/bijectors_utils.jl

@@ -199,6 +199,7 @@ test_with_flux = (scenario) -> begin
             scenf = (scenario..., :use_Flux, :use_gpu, :omit_r0)
             rng = StableRNG(111)
             # here using DoubleMMCase() directly rather than construct_problem
+            #(;transP, transM) = get_hybridproblem_transforms(DoubleMM.DoubleMMCase(); scenario = scenf)
             prob = probg = HybridProblem(DoubleMM.DoubleMMCase(); scenario = scenf);
             solver = HybridPosteriorSolver(; alg=Adam(0.02),  n_MC=3)
             n_site, n_batch = get_hybridproblem_n_site_and_batch(prob; scenario = scenf)
@@ -216,13 +217,15 @@ test_with_flux = (scenario) -> begin
                 maxiters = 37, 
             );
             @test cdev(ϕ.unc.ρsM)[1] > 0 
+            @test probo.ϕunc == cdev(ϕ.unc)
             test_correlation = () -> begin
-                n_epoch = 100 # requires 
+                n_epoch = 20 # requires 
                 (; ϕ, θP, resopt, probo) = solve(prob, solver; scenario = scenf,
                     maxiters = n_batches_in_epoch * n_epoch, 
                     callback = callback_loss(n_batches_in_epoch*5)
                 );
                 @test cdev(ϕ.unc.ρsM)[1] > 0 
+                @test probo.ϕunc == cdev(ϕ.unc)
                 # predict using problem and its associated dataloader
                 (; θ, y, entropy_ζ) = predict_gf(rng, probo; scenario = scenf, n_sample_pred = 200);            
                 mean_θ = CA.ComponentVector(mean(CA.getdata(θ); dims = 2)[:, 1], CA.getaxes(θ[:, 1])[1])
@@ -240,8 +243,7 @@ test_with_flux = (scenario) -> begin
             scenf = (scenario..., :use_Flux, :use_gpu, :omit_r0, :f_on_gpu)
             rng = StableRNG(111)
             probg = HybridProblem(DoubleMM.DoubleMMCase(); scenario = scenf);
-            # TODO replace exp by Exp Transformer
-            prob = CP.update(probg, transM = identity, transP = identity);
+            #prob = CP.update(probg, transM = identity, transP = identity);
             solver = HybridPosteriorSolver(; alg=Adam(0.02), n_MC=3)
             n_site, n_batch = get_hybridproblem_n_site_and_batch(prob; scenario = scenf)
             n_batches_in_epoch = n_site ÷ n_batch

test/test_HybridProblem.jl

@@ -22,6 +22,7 @@ b2 = elementwise(exp)
 b2s = Stacked((b2,b2),(1:3,4:4))
 b3 = HybridVariationalInference.Exp()
 b3s = Stacked((b3,b3), (1:3,4:4))
+#b3s = Stacked((b3,),(1:4,))
 
 
 y = trans(x, b2)
@@ -35,6 +36,10 @@ dy = Zygote.gradient(x -> trans(x,b2), x)
 end;
 
 @testset "Exp" begin
+    y1 = b3(x)
+    y2 = b3s(x)
+    @test all(inverse(b3)(y2) .≈ x)
+    @test all(inverse(b3s)(y2) .≈ x)
     ye = trans(x, b3)
     dye = Zygote.gradient(x -> trans(x,b3), x)
     @test ye == y