enable variance reduction by default for

Author: BatyLeo

src/perturbed/abstract_perturbed.jl

@@ -97,7 +97,7 @@ end
 Forward pass. Compute the expectation of the underlying distribution.
 """
 function (perturbed::AbstractPerturbed)(
-    θ::AbstractArray; autodiff_variance_reduction::Bool=false, kwargs...
+    θ::AbstractArray; autodiff_variance_reduction::Bool=true, kwargs...
 )
     probadist = compute_probability_distribution(
         perturbed, θ; autodiff_variance_reduction, kwargs...
@@ -118,7 +118,7 @@ function ChainRulesCore.rrule(
     ::typeof(compute_probability_distribution),
     perturbed::AbstractPerturbed,
     θ::AbstractArray;
-    autodiff_variance_reduction::Bool=false,
+    autodiff_variance_reduction::Bool=true,
     kwargs...,
 )
     η_samples = sample_perturbations(perturbed, θ)

test/jacobian_approx.jl

@@ -15,8 +15,10 @@
 
     @testset "PerturbedAdditive" begin
         # Compute jacobian with reverse mode
-        jac1 = Zygote.jacobian(perturbed1, θ)[1]
-        jac1_big = Zygote.jacobian(perturbed1_big, θ)[1]
+        jac1 = Zygote.jacobian(θ -> perturbed1(θ; autodiff_variance_reduction=false), θ)[1]
+        jac1_big = Zygote.jacobian(
+            θ -> perturbed1_big(θ; autodiff_variance_reduction=false), θ
+        )[1]
         # Only diagonal should be positive
         @test all(diag(jac1) .>= 0)
         @test all(jac1 - Diagonal(jac1) .<= 0)
@@ -27,8 +29,10 @@
     end
 
     @testset "PerturbedMultiplicative" begin
-        jac2 = Zygote.jacobian(perturbed2, θ)[1]
-        jac2_big = Zygote.jacobian(perturbed2_big, θ)[1]
+        jac2 = Zygote.jacobian(θ -> perturbed2(θ; autodiff_variance_reduction=false), θ)[1]
+        jac2_big = Zygote.jacobian(
+            θ -> perturbed2_big(θ; autodiff_variance_reduction=false), θ
+        )[1]
         @test all(diag(jac2) .>= 0)
         @test all(jac2 - Diagonal(jac2) .<= 0)
         @test sortperm(diag(jac2)) != sortperm(θ)

test/perturbed_oracle.jl

@@ -31,11 +31,11 @@ end
     n = 10
     θ = randn(10)
 
-    Ja = jacobian(pa, θ)[1]
-    Ja_reduced_variance = jacobian(x -> pa(x; autodiff_variance_reduction=true), θ)[1]
+    Ja = jacobian(θ -> pa(θ; autodiff_variance_reduction=false), θ)[1]
+    Ja_reduced_variance = jacobian(pa, θ)[1]
 
-    Jm = jacobian(pm, θ)[1]
-    Jm_reduced_variance = jacobian(x -> pm(x; autodiff_variance_reduction=true), θ)[1]
+    Jm = jacobian(x -> pm(x; autodiff_variance_reduction=false), θ)[1]
+    Jm_reduced_variance = jacobian(pm, θ)[1]
 
     J_true = Matrix(I, n, n)  # exact jacobian is the identity matrix