add elbo_batch implementation; much faster for invertible NN based flows

Author: zuhengxu

src/NormalizingFlows.jl

@@ -12,7 +12,7 @@ import DifferentiationInterface as DI
 
 using DocStringExtensions
 
-export train_flow, elbo, loglikelihood
+export train_flow, elbo, elbo_batch, loglikelihood
 
 """
     train_flow([rng::AbstractRNG, ]vo, flow, args...; kwargs...)

src/objectives/elbo.jl

@@ -43,3 +43,39 @@ end
 function elbo(flow::Bijectors.TransformedDistribution, logp, n_samples)
     return elbo(Random.default_rng(), flow, logp, n_samples)
 end
+
+
+"""
+    elbo_batch(flow, logp, xs) 
+    elbo_batch([rng, ]flow, logp, n_samples)
+
+Instead of broadcasting over elbo_single_sample, this function directly
+computes the ELBO in a batched manner, which requires the flow.transform to be able to 
+handle batched transformation directly.
+
+This will be more efficient than `elbo` for invertible neural networks such as RealNVP, 
+Neural Spline Flow, etc.
+
+# Arguments
+- `rng`: random number generator
+- `flow`: variational distribution to be trained. In particular 
+  `flow = transformed(q₀, T::Bijectors.Bijector)`, 
+  q₀ is a reference distribution that one can easily sample and compute logpdf
+- `logp`: log-pdf of the target distribution (not necessarily normalized)
+- `xs`: samples from reference dist q₀
+- `n_samples`: number of samples from reference dist q₀
+
+"""
+function elbo_batch(flow::Bijectors.MultivariateTransformed, logp, xs::AbstractMatrix)
+    # requires the flow transformation to be able to handle batched inputs
+    ys, logabsdetjac = with_logabsdet_jacobian(flow.transform, xs) 
+    elbos = logp(ys) .- logpdf(flow.dist, xs) .+ logabsdetjac
+    return elbos
+end
+function elbo_batch(rng::AbstractRNG, flow::Bijectors.MultivariateTransformed, logp, n_samples)
+    xs = _device_specific_rand(rng, flow.dist, n_samples)
+    elbos = elbo_batch(flow, logp, xs)
+    return mean(elbos)
+end
+elbo_batch(flow::Bijectors.UnivariateTransformed, logp, n_samples) = 
+    elbo_batch(Random.default_rng(), flow, logp, n_samples)