Perfect coverage

README.md

@@ -17,3 +17,5 @@ The following estimators are implemented:
 
   - [REINFORCE](https://jmlr.org/papers/volume21/19-346/19-346.pdf#section.20)
   - [Reparametrization](https://jmlr.org/papers/volume21/19-346/19-346.pdf#section.56)
+
+> Warning: this package is experimental, use at your own risk and expect frequent breaking releases.

test/expectation.jl

@@ -11,6 +11,8 @@ using Zygote
 exp_with_kwargs(x; correct=false) = correct ? exp(x) : sin(x)
 vec_exp_with_kwargs(x; correct=false) = exp_with_kwargs.(x; correct)
 
+normal_logdensity_grad(x, θ...) = gradient((_θ...) -> logpdf(Normal(_θ...), x), θ...)
+
 μ, σ = 0.5, 1.0
 true_mean(μ, σ) = mean(LogNormal(μ, σ))
 true_std(μ, σ) = std(LogNormal(μ, σ))
@@ -26,6 +28,14 @@ true_std(μ, σ) = std(LogNormal(μ, σ))
                 nb_samples=10^4,
                 threaded=threaded,
             ),
+            Reinforce(
+                exp_with_kwargs,
+                Normal,
+                normal_logdensity_grad;
+                rng=StableRNG(63),
+                nb_samples=10^4,
+                threaded=threaded,
+            ),
             Reparametrization(
                 exp_with_kwargs,
                 Normal;