fixed error in doc for creat_flow

Author: zuhengxu

docs/src/example.md

@@ -9,7 +9,7 @@ attaching them to the base distribution `q0`.
 
 - `layers`: an iterable of `Bijectors.Bijector` objects that are composed in order
   (left-to-right) via function composition 
-(for instance, if `layers = [l1, l2, l3]`, the flow will be `l3∘l2∘l1(q0)`).
+(for instance, if `layers = [l1, l2, l3]`, the flow will be `l1∘l2∘l3(q0)`).
 - `q0`: the base distribution (e.g., `MvNormal(zeros(d), I)`).
 
 Returns a `Bijectors.TransformedDistribution` representing the resulting flow.
@@ -18,7 +18,7 @@ Example
 
     using Distributions, Bijectors, LinearAlgebra
     q0 = MvNormal(zeros(2), I)
-    flow = create_flow((Bijectors.Scale([1.0, 2.0]), Bijectors.Shift([0.0, 1.0])), q0)
+    flow = create_flow((Bijectors.Shift([0.0, 1.0]), Bijectors.Scale([1.0, 2.0])), q0)
 """
 function create_flow(Ls, q₀)
     ts =  reduce(∘, Ls)
@@ -62,7 +62,7 @@ Create a fully connected neural network (FNN).
 - `hidden_dims::AbstractVector{<:Int}`: A vector of integers specifying the dimensions of the hidden layers.
 - `output_dim::Int`: The dimension of the output layer.
 - `inlayer_activation`: The activation function for the hidden layers. Defaults to `Flux.leakyrelu`.
-- `output_activation`: The activation function for the output layer. Defaults to `Flux.tanh`.
+- `output_activation`: The activation function for the output layer. Defaults to `nothing`.
 - `paramtype::Type{T} = Float64`: The type of the parameters in the network, defaults to `Float64`.
 
 # Returns

src/flows/utils.jl

@@ -46,6 +46,29 @@ function elbo(flow::Bijectors.TransformedDistribution, logp, n_samples)
 end
 
 
+"""
+    _batched_elbos(flow, logp, xs)
+
+Batched ELBO estimates that transforms a matrix of samples (each column represents a single
+sample) in one call. 
+This is more efficient for invertible neural-network flows (RealNVP/NSF) as it leverages
+the batched operation of the neural networks.
+
+Inputs
+- `flow::Bijectors.MultivariateTransformed`
+- `logp`: function returning log-density of target
+- `xs`: column-wise sample batch
+
+Returns
+- a vector of ELBO estimates for each sample in the batch
+"""
+function _batched_elbos(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
+
 """
     elbo_batch(flow, logp, xs)
     elbo_batch([rng, ] flow, logp, n_samples)
@@ -64,14 +87,12 @@ Returns
 - Scalar estimate of the ELBO
 """
 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
+    elbos = _batched_elbos(flow, logp, xs)
+    return mean(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)
+    elbos = _batched_elbos(flow, logp, xs)
     return mean(elbos)
 end
 elbo_batch(flow::Bijectors.TransformedDistribution, logp, n_samples) =