use elbo_batch for real_nvp; achieved 4+ times speed up

Author: zuhengxu

example/demo_RealNVP.jl

@@ -36,7 +36,7 @@ function AffineCoupling(
     return AffineCoupling(dim, mask, s, t)
 end
 
-function Bijectors.transform(af::AffineCoupling, x::AbstractVector)
+function Bijectors.transform(af::AffineCoupling, x::AbstractVecOrMat)
     # partition vector using 'af.mask::PartitionMask`
     x₁, x₂, x₃ = partition(af.mask, x)
     y₁ = x₁ .* af.s(x₂) .+ af.t(x₂)
@@ -50,10 +50,18 @@ end
 function Bijectors.with_logabsdet_jacobian(af::AffineCoupling, x::AbstractVector)
     x_1, x_2, x_3 = Bijectors.partition(af.mask, x)
     y_1 = af.s(x_2) .* x_1 .+ af.t(x_2)
-    logjac = sum(log ∘ abs, af.s(x_2))
+    logjac = sum(log ∘ abs, af.s(x_2)) # this is a scalar
     return combine(af.mask, y_1, x_2, x_3), logjac
 end
 
+function Bijectors.with_logabsdet_jacobian(af::AffineCoupling, x::AbstractMatrix)
+    x_1, x_2, x_3 = Bijectors.partition(af.mask, x)
+    y_1 = af.s(x_2) .* x_1 .+ af.t(x_2)
+    logjac = sum(log ∘ abs, af.s(x_2); dims = 1) # 1 × size(x, 2)
+    return combine(af.mask, y_1, x_2, x_3), vec(logjac)
+end
+
+
 function Bijectors.with_logabsdet_jacobian(
     iaf::Inverse{<:AffineCoupling}, y::AbstractVector
 )
@@ -66,10 +74,16 @@ function Bijectors.with_logabsdet_jacobian(
     return combine(af.mask, x_1, y_2, y_3), logjac
 end
 
-function Bijectors.logabsdetjac(af::AffineCoupling, x::AbstractVector)
-    _, x_2, _ = partition(af.mask, x)
-    logjac = sum(log ∘ abs, af.s(x_2))
-    return logjac
+function Bijectors.with_logabsdet_jacobian(
+    iaf::Inverse{<:AffineCoupling}, y::AbstractMatrix
+)
+    af = iaf.orig
+    # partition vector using `af.mask::PartitionMask`
+    y_1, y_2, y_3 = partition(af.mask, y)
+    # inverse transformation
+    x_1 = (y_1 .- af.t(y_2)) ./ af.s(y_2)
+    logjac = -sum(log ∘ abs, af.s(y_2); dims = 1)
+    return combine(af.mask, x_1, y_2, y_3), vec(logjac)
 end
 
 ################### 
@@ -126,6 +140,8 @@ q0 = MvNormal(zeros(T, 2), ones(T, 2))
 
 d = 2
 hdims = 32
+
+# alternating the coupling layers
 Ls = [AffineCoupling(d, hdims, [1]) ∘ AffineCoupling(d, hdims, [2]) for i in 1:3]
 
 flow = create_flow(Ls, q0)
@@ -142,7 +158,8 @@ cb(iter, opt_stats, re, θ) = (sample_per_iter=sample_per_iter,ad=adtype)
 adtype = ADTypes.AutoMooncake(; config = Mooncake.Config())
 checkconv(iter, stat, re, θ, st) = stat.gradient_norm < one(T)/1000
 flow_trained, stats, _ = train_flow(
-    elbo,
+    rng, 
+    elbo_batch,        # using elbo_batch instead of elbo achieves 4-5 times speedup 
     flow,
     logp,
     sample_per_iter;