minor ed on docs

Author: zuhengxu

docs/src/api.md

@@ -63,8 +63,8 @@ By maximizing the ELBO, it is equivalent to minimizing the reverse KL divergence
 \begin{aligned}
 &\min _{\theta} \mathbb{E}_{q_{\theta}}\left[\log q_{\theta}(Z)-\log p(Z)\right]  \quad \text{(Reverse KL)}\\
 & = \max _{\theta} \mathbb{E}_{q_0}\left[ \log p\left(T_N \circ \cdots \circ
-T_1(Z_0)\right)-\log q_0(X)+\sum_{n=1}^N \log J_n\left(F_n \circ \cdots \circ
-F_1(X)\right)\right] \quad \text{(ELBO)}
+T_1(Z_0)\right)-\log q_0(X)+\sum_{n=1}^N \log J_n\left(T_n \circ \cdots \circ
+T_1(X)\right)\right] \quad \text{(ELBO)}
 \end{aligned}
 ```
 

example/demo_neural_spline_flow.jl

@@ -48,7 +48,7 @@ flow_trained, stats, _ = train_flow(
     flow,
     logp,
     sample_per_iter;
-    max_iters=10000,   # change to larger number of iterations (e.g., 50_000) for better results
+    max_iters=10,   # change to larger number of iterations (e.g., 50_000) for better results
     optimiser=Optimisers.Adam(1e-4),
     ADbackend=adtype,
     show_progress=true,

src/objectives/elbo.jl

@@ -50,8 +50,9 @@ end
     elbo_batch(flow, logp, xs)
     elbo_batch([rng, ] flow, logp, n_samples)
 
-Batched ELBO evaluation that transforms a matrix of samples in one call. This
-is more efficient for invertible neural-network flows (RealNVP/NSF) as it leverages
+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