From 99d293fc724ccf21305238dd77259edcd827b79f Mon Sep 17 00:00:00 2001
From: Seth Axen <seth.axen@gmail.com>
Date: Mon, 25 Apr 2022 00:08:03 +0200
Subject: [PATCH 01/23] Add Transducers and Folds as dependencies

---
 Project.toml       |  4 ++++
 src/Pathfinder.jl  |  3 +++
 src/transducers.jl | 27 +++++++++++++++++++++++++++
 3 files changed, 34 insertions(+)
 create mode 100644 src/transducers.jl

diff --git a/Project.toml b/Project.toml
index 63d709ec..00a80f51 100644
--- a/Project.toml
+++ b/Project.toml
@@ -6,6 +6,7 @@ version = "0.3.5"
 [deps]
 AbstractDifferentiation = "c29ec348-61ec-40c8-8164-b8c60e9d9f3d"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
+Folds = "41a02a25-b8f0-4f67-bc48-60067656b558"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 GalacticOptim = "a75be94c-b780-496d-a8a9-0878b188d577"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
@@ -17,10 +18,12 @@ Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
+Transducers = "28d57a85-8fef-5791-bfe6-a80928e7c999"
 
 [compat]
 AbstractDifferentiation = "0.4"
 Distributions = "0.25"
+Folds = "0.2"
 ForwardDiff = "0.10"
 GalacticOptim = "2"
 Optim = "1.4"
@@ -29,6 +32,7 @@ PSIS = "0.2, 0.3, 0.4"
 Requires = "1"
 StatsBase = "0.33"
 StatsFuns = "0.9"
+Transducers = "0.4.5"
 julia = "1.6"
 
 [extras]
diff --git a/src/Pathfinder.jl b/src/Pathfinder.jl
index 745f6af0..dafde4a5 100644
--- a/src/Pathfinder.jl
+++ b/src/Pathfinder.jl
@@ -2,6 +2,7 @@ module Pathfinder
 
 using AbstractDifferentiation: AD
 using Distributions: Distributions
+using Folds: Folds
 # ensure that ForwardDiff is conditionally loaded by GalacticOptim
 using ForwardDiff: ForwardDiff
 using LinearAlgebra
@@ -14,6 +15,7 @@ using Requires: Requires
 using Statistics: Statistics
 using StatsBase: StatsBase
 using StatsFuns: log2π
+using Transducers: Transducers
 
 export pathfinder, multipathfinder
 
@@ -25,6 +27,7 @@ const DEFAULT_OPTIMIZER = Optim.LBFGS(;
     m=DEFAULT_HISTORY_LENGTH, linesearch=DEFAULT_LINE_SEARCH
 )
 
+include("transducers.jl")
 include("woodbury.jl")
 include("optimize.jl")
 include("inverse_hessian.jl")
diff --git a/src/transducers.jl b/src/transducers.jl
new file mode 100644
index 00000000..19e97457
--- /dev/null
+++ b/src/transducers.jl
@@ -0,0 +1,27 @@
+# utility functions to work with Transducers.jl
+
+_is_thread_safe_rng(rng) = false
+if VERSION ≥ v"1.7.0"
+    _is_thread_safe_rng(rng::Random.TaskLocalRNG) = true
+    _is_thread_safe_rng(rng::typeof(Random.GLOBAL_RNG)) = true
+end
+
+function _default_executor(rng; kwargs...)
+    if _is_thread_safe_rng(rng)
+        return Transducers.PreferParallel(; kwargs...)
+    else
+        return Transducers.SequentialEx()
+    end
+end
+
+# transducer-friendly findmax
+function _findmax(x)
+    return Transducers.foldxl(x |> Transducers.Enumerate(); init=missing) do xmax_imax, i_xi
+        xmax_imax === missing && return reverse(i_xi)
+        return last(i_xi) > first(xmax_imax) ? reverse(i_xi) : xmax_imax
+    end
+end
+
+# WARNING: Type piracy!
+# https://github.com/JuliaFolds/Transducers.jl/issues/521
+Base.size(x::Transducers.ProgressLoggingFoldable) = size(x.foldable)

From d5c3827ff2c23b0c4fd1794cee6bfe27125a2f54 Mon Sep 17 00:00:00 2001
From: Seth Axen <seth.axen@gmail.com>
Date: Mon, 25 Apr 2022 00:09:28 +0200
Subject: [PATCH 02/23] Take executor for running paths in parallel

---
 src/multipath.jl | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/src/multipath.jl b/src/multipath.jl
index 37c8d019..f5f0b18c 100644
--- a/src/multipath.jl
+++ b/src/multipath.jl
@@ -37,6 +37,10 @@ resulting draws better approximate draws from the target distribution ``p`` inst
 - `ad_backend=AD.ForwardDiffBackend()`: AbstractDifferentiation.jl AD backend.
 - `ndraws_per_run::Int=5`: The number of draws to take for each component before resampling.
 - `importance::Bool=true`: Perform Pareto smoothed importance resampling of draws.
+- `executor::Transducers.Executor`: Transducers.jl executor that determines if and how
+    to run the single-path runs in parallel. If `rng` is known to be thread-safe, the
+    default is `Transducers.PreferParallel(; basesize=1)` (parallel executation, defaulting
+    to multi-threading). Otherwise, it is `Transducers.SerialEx()` (no parallelization).
 - `kwargs...` : Remaining keywords are forwarded to [`pathfinder`](@ref).
 
 # Returns
@@ -81,6 +85,7 @@ function multipathfinder(
     ndraws;
     ndraws_per_run::Int=5,
     rng::Random.AbstractRNG=Random.default_rng(),
+    executor::Transducers.Executor=_default_executor(rng; basesize=1),
     importance::Bool=true,
     kwargs...,
 )
@@ -94,11 +99,14 @@ function multipathfinder(
 
     # run pathfinder independently from each starting point
     # TODO: allow to be parallelized
-    res = map(θ₀s) do θ₀
-        return pathfinder(optim_fun, θ₀, ndraws_per_run; rng, kwargs...)
+    trans = Transducers.Map() do θ₀
+        return pathfinder(
+            optim_fun, θ₀, ndraws_per_run; rng, executor=executor_per_run, kwargs...
+        )
     end
     qs = reduce(vcat, first.(res))
     ϕs = reduce(hcat, getindex.(res, 2))
+    res = Folds.collect(iter_sp, executor)
 
     # draw samples from augmented mixture model
     inds = axes(ϕs, 2)

From bf4b9e64703dee0e0cc6bdd00be896ac190981ec Mon Sep 17 00:00:00 2001
From: Seth Axen <seth.axen@gmail.com>
Date: Mon, 25 Apr 2022 00:10:16 +0200
Subject: [PATCH 03/23] Take executor for computing elbo in parallel

---
 src/elbo.jl       | 7 ++++---
 src/singlepath.jl | 8 +++++++-
 2 files changed, 11 insertions(+), 4 deletions(-)

diff --git a/src/elbo.jl b/src/elbo.jl
index c9962076..b051f0e2 100644
--- a/src/elbo.jl
+++ b/src/elbo.jl
@@ -1,8 +1,9 @@
-function maximize_elbo(rng, logp, dists, ndraws)
-    elbo_ϕ_logqϕ = map(dists) do dist
+function maximize_elbo(rng, logp, dists, ndraws, executor)
+    iter = dists |> Transducers.Map() do dist
         elbo_and_samples(rng, logp, dist, ndraws)
     end
-    lopt = argmax(first.(elbo_ϕ_logqϕ))
+    elbo_ϕ_logqϕ = Folds.collect(iter, executor)
+    _, lopt = _findmax(elbo_ϕ_logqϕ |> Transducers.Map(first))
     return (lopt, elbo_ϕ_logqϕ[lopt]...)
 end
 
diff --git a/src/singlepath.jl b/src/singlepath.jl
index f910ad3b..9f582a61 100644
--- a/src/singlepath.jl
+++ b/src/singlepath.jl
@@ -21,6 +21,10 @@ constructed using at most the previous `history_length` steps.
 # Keywords
 - `ad_backend=AD.ForwardDiffBackend()`: AbstractDifferentiation.jl AD backend.
 - `rng::Random.AbstractRNG`: The random number generator to be used for drawing samples
+- `executor::Transducers.Executor`: Transducers.jl executor that determines if and how
+    to perform ELBO computation in parallel. If `rng` is known to be thread-safe, the
+    default is `Transducers.PreferParallel()` (parallel executation, defaulting to
+    multi-threading). Otherwise, it is `Transducers.SerialEx()` (no parallelization).
 - `optimizer`: Optimizer to be used for constructing trajectory. Can be any optimizer
     compatible with GalacticOptim, so long as it supports callbacks. Defaults to
     `Optim.LBFGS(; m=$DEFAULT_HISTORY_LENGTH, linesearch=LineSearches.MoreThuente())`. See
@@ -66,6 +70,7 @@ function pathfinder(
     θ₀,
     ndraws;
     rng::Random.AbstractRNG=Random.default_rng(),
+    executor::Transducers.Executor=_default_executor(rng),
     optimizer=DEFAULT_OPTIMIZER,
     history_length::Int=optimizer isa Optim.LBFGS ? optimizer.m : DEFAULT_HISTORY_LENGTH,
     ndraws_elbo::Int=5,
@@ -92,6 +97,7 @@ function pathfinder(
     optim_prob::GalacticOptim.OptimizationProblem,
     ndraws;
     rng::Random.AbstractRNG=Random.default_rng(),
+    executor::Transducers.Executor=_default_executor(rng),
     optimizer=DEFAULT_OPTIMIZER,
     history_length::Int=optimizer isa Optim.LBFGS ? optimizer.m : DEFAULT_HISTORY_LENGTH,
     ndraws_elbo::Int=5,
@@ -109,7 +115,7 @@ function pathfinder(
     qs = fit_mvnormals(θs, ∇logpθs; history_length)
 
     # find ELBO-maximizing distribution
-    lopt, elbo, ϕ, logqϕ = maximize_elbo(rng, logp, qs[2:end], ndraws_elbo)
+    lopt, elbo, ϕ, logqϕ = maximize_elbo(rng, logp, qs[2:end], ndraws_elbo, executor)
     @info "Optimized for $L iterations. Maximum ELBO of $(round(elbo; digits=2)) reached at iteration $lopt."
 
     # get parameters of ELBO-maximizing distribution

From 32aae721a53f10c8765652792baa48d186ecb645 Mon Sep 17 00:00:00 2001
From: Seth Axen <seth.axen@gmail.com>
Date: Mon, 25 Apr 2022 00:10:34 +0200
Subject: [PATCH 04/23] Allow privision of per-run executor

---
 src/multipath.jl | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/src/multipath.jl b/src/multipath.jl
index f5f0b18c..326da598 100644
--- a/src/multipath.jl
+++ b/src/multipath.jl
@@ -41,6 +41,8 @@ resulting draws better approximate draws from the target distribution ``p`` inst
     to run the single-path runs in parallel. If `rng` is known to be thread-safe, the
     default is `Transducers.PreferParallel(; basesize=1)` (parallel executation, defaulting
     to multi-threading). Otherwise, it is `Transducers.SerialEx()` (no parallelization).
+- `executor_per_run::Transducers.Executor`: Transducers.jl executor used within each run to
+    parallelize PRNG calls. See [`pathfinder`](@ref) for a description.
 - `kwargs...` : Remaining keywords are forwarded to [`pathfinder`](@ref).
 
 # Returns
@@ -86,6 +88,7 @@ function multipathfinder(
     ndraws_per_run::Int=5,
     rng::Random.AbstractRNG=Random.default_rng(),
     executor::Transducers.Executor=_default_executor(rng; basesize=1),
+    executor_per_run=_default_executor(rng),
     importance::Bool=true,
     kwargs...,
 )

From 2f109a6b6893ce6f3e5a4f4388c54dc0361a7438 Mon Sep 17 00:00:00 2001
From: Seth Axen <seth.axen@gmail.com>
Date: Mon, 25 Apr 2022 00:11:17 +0200
Subject: [PATCH 05/23] Default to GLOBAL_RNG

---
 src/multipath.jl  | 4 +++-
 src/singlepath.jl | 4 ++--
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/src/multipath.jl b/src/multipath.jl
index 326da598..76835c74 100644
--- a/src/multipath.jl
+++ b/src/multipath.jl
@@ -37,6 +37,8 @@ resulting draws better approximate draws from the target distribution ``p`` inst
 - `ad_backend=AD.ForwardDiffBackend()`: AbstractDifferentiation.jl AD backend.
 - `ndraws_per_run::Int=5`: The number of draws to take for each component before resampling.
 - `importance::Bool=true`: Perform Pareto smoothed importance resampling of draws.
+- `rng::AbstractRNG=Random.GLOBAL_RNG`: Pseudorandom number generator. It is recommended to
+    use a parallelization-friendly PRNG like the default PRNG on Julia 1.7 and up.
 - `executor::Transducers.Executor`: Transducers.jl executor that determines if and how
     to run the single-path runs in parallel. If `rng` is known to be thread-safe, the
     default is `Transducers.PreferParallel(; basesize=1)` (parallel executation, defaulting
@@ -86,7 +88,7 @@ function multipathfinder(
     θ₀s,
     ndraws;
     ndraws_per_run::Int=5,
-    rng::Random.AbstractRNG=Random.default_rng(),
+    rng::Random.AbstractRNG=Random.GLOBAL_RNG,
     executor::Transducers.Executor=_default_executor(rng; basesize=1),
     executor_per_run=_default_executor(rng),
     importance::Bool=true,
diff --git a/src/singlepath.jl b/src/singlepath.jl
index 9f582a61..8431199b 100644
--- a/src/singlepath.jl
+++ b/src/singlepath.jl
@@ -69,7 +69,7 @@ function pathfinder(
     optim_fun::GalacticOptim.OptimizationFunction,
     θ₀,
     ndraws;
-    rng::Random.AbstractRNG=Random.default_rng(),
+    rng::Random.AbstractRNG=Random.GLOBAL_RNG,
     executor::Transducers.Executor=_default_executor(rng),
     optimizer=DEFAULT_OPTIMIZER,
     history_length::Int=optimizer isa Optim.LBFGS ? optimizer.m : DEFAULT_HISTORY_LENGTH,
@@ -96,7 +96,7 @@ See [`pathfinder`](@ref) for a description of remaining arguments.
 function pathfinder(
     optim_prob::GalacticOptim.OptimizationProblem,
     ndraws;
-    rng::Random.AbstractRNG=Random.default_rng(),
+    rng::Random.AbstractRNG=Random.GLOBAL_RNG,
     executor::Transducers.Executor=_default_executor(rng),
     optimizer=DEFAULT_OPTIMIZER,
     history_length::Int=optimizer isa Optim.LBFGS ? optimizer.m : DEFAULT_HISTORY_LENGTH,

From ab37cdae965e3454b08454abcf842e3027f321f1 Mon Sep 17 00:00:00 2001
From: Seth Axen <seth.axen@gmail.com>
Date: Mon, 25 Apr 2022 00:11:51 +0200
Subject: [PATCH 06/23] Forward executor

---
 src/singlepath.jl | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/singlepath.jl b/src/singlepath.jl
index 8431199b..a01b1e12 100644
--- a/src/singlepath.jl
+++ b/src/singlepath.jl
@@ -77,7 +77,9 @@ function pathfinder(
     kwargs...,
 )
     optim_prob = build_optim_problem(optim_fun, θ₀; kwargs...)
-    return pathfinder(optim_prob, ndraws; rng, optimizer, history_length, ndraws_elbo)
+    return pathfinder(
+        optim_prob, ndraws; rng, executor, optimizer, history_length, ndraws_elbo
+    )
 end
 
 """

From 0ab78cfc03048081faf6e2b09a14d49360b36447 Mon Sep 17 00:00:00 2001
From: Seth Axen <seth.axen@gmail.com>
Date: Mon, 25 Apr 2022 00:13:35 +0200
Subject: [PATCH 07/23] Add missing iterator

---
 src/multipath.jl | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/multipath.jl b/src/multipath.jl
index 76835c74..5f4415f8 100644
--- a/src/multipath.jl
+++ b/src/multipath.jl
@@ -111,6 +111,7 @@ function multipathfinder(
     end
     qs = reduce(vcat, first.(res))
     ϕs = reduce(hcat, getindex.(res, 2))
+    iter_sp = Transducers.withprogress(θ₀s; interval=1e-3) |> trans
     res = Folds.collect(iter_sp, executor)
 
     # draw samples from augmented mixture model

From 73a768cfeda8ed22d06cdc9f79df03a894855b3a Mon Sep 17 00:00:00 2001
From: Seth Axen <seth.axen@gmail.com>
Date: Mon, 25 Apr 2022 00:14:18 +0200
Subject: [PATCH 08/23] Use more transducers

---
 src/elbo.jl      |  2 +-
 src/multipath.jl | 10 ++++++----
 src/mvnormal.jl  |  7 +++++--
 src/optimize.jl  |  2 +-
 4 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/src/elbo.jl b/src/elbo.jl
index b051f0e2..9b86abb6 100644
--- a/src/elbo.jl
+++ b/src/elbo.jl
@@ -9,7 +9,7 @@ end
 
 function elbo_and_samples(rng, logp, dist, ndraws)
     ϕ, logqϕ = rand_and_logpdf(rng, dist, ndraws)
-    logpϕ = logp.(eachcol(ϕ))
+    logpϕ = eachcol(ϕ) |> Transducers.Map(logp) |> Transducers.tcollect
     elbo = elbo_from_logpdfs(logpϕ, logqϕ)
     return elbo, ϕ, logqϕ
 end
diff --git a/src/multipath.jl b/src/multipath.jl
index 5f4415f8..2c2f28fc 100644
--- a/src/multipath.jl
+++ b/src/multipath.jl
@@ -109,16 +109,18 @@ function multipathfinder(
             optim_fun, θ₀, ndraws_per_run; rng, executor=executor_per_run, kwargs...
         )
     end
-    qs = reduce(vcat, first.(res))
-    ϕs = reduce(hcat, getindex.(res, 2))
     iter_sp = Transducers.withprogress(θ₀s; interval=1e-3) |> trans
     res = Folds.collect(iter_sp, executor)
+    qs = res |> Transducers.Map(first) |> collect
+    ϕs = reduce(hcat, res |> Transducers.Map(x -> x[2]))
 
     # draw samples from augmented mixture model
     inds = axes(ϕs, 2)
     sample_inds = if importance
-        logqϕs = reduce(vcat, last.(res))
-        log_ratios = map(((ϕ, logqϕ),) -> logp(ϕ) - logqϕ, zip(eachcol(ϕs), logqϕs))
+        logqϕs = res |> Transducers.MapCat(last) |> collect
+        iter_logp = eachcol(ϕs) |> Transducers.Map(logp)
+        logpϕs = Folds.collect(iter_logp, executor)
+        log_ratios = logpϕs - logqϕs
         resample(rng, inds, log_ratios, ndraws)
     else
         resample(rng, inds, ndraws)
diff --git a/src/mvnormal.jl b/src/mvnormal.jl
index e1232593..9d45d983 100644
--- a/src/mvnormal.jl
+++ b/src/mvnormal.jl
@@ -10,9 +10,12 @@ point are returned.
 """
 function fit_mvnormals(θs, ∇logpθs; kwargs...)
     Σs = lbfgs_inverse_hessians(θs, ∇logpθs; kwargs...)
+    trans = Transducers.MapSplat() do Σ, ∇logpθ, θ
+        μ = muladd(Σ, ∇logpθ, θ)
+        return Distributions.MvNormal(μ, Σ)
+    end
     l = length(Σs)
-    μs = @views map(muladd, Σs, ∇logpθs[1:l], θs[1:l])
-    return Distributions.MvNormal.(μs, Σs)
+    return @views(zip(Σs, ∇logpθs[1:l], θs[1:l])) |> trans |> collect
 end
 
 # faster than computing `logpdf` and `rand` independently
diff --git a/src/optimize.jl b/src/optimize.jl
index 96746b1a..d901a077 100644
--- a/src/optimize.jl
+++ b/src/optimize.jl
@@ -52,6 +52,6 @@ function optimize_with_trace(prob, optimizer)
     # NOTE: GalacticOptim doesn't have an interface for accessing the gradient trace,
     # so we need to recompute it ourselves
     # see https://github.com/SciML/GalacticOptim.jl/issues/149
-    ∇fxs = ∇f.(xs)
+    ∇fxs = xs |> Transducers.Map(∇f) |> Transducers.tcollect
     return xs, fxs, ∇fxs
 end

From 3ee7cd985a6f1b712abb5e4f4ac500e1270e9456 Mon Sep 17 00:00:00 2001
From: Seth Axen <seth.axen@gmail.com>
Date: Mon, 25 Apr 2022 00:14:38 +0200
Subject: [PATCH 09/23] Allow pipe syntax for transducers

---
 .JuliaFormatter.toml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/.JuliaFormatter.toml b/.JuliaFormatter.toml
index 323237ba..5f4dff16 100644
--- a/.JuliaFormatter.toml
+++ b/.JuliaFormatter.toml
@@ -1 +1,2 @@
 style = "blue"
+pipe_to_function_call = false

From cd23bea6d93efc40c22fbef4e0f53990aa5f18e1 Mon Sep 17 00:00:00 2001
From: Seth Axen <seth.axen@gmail.com>
Date: Mon, 25 Apr 2022 00:17:49 +0200
Subject: [PATCH 10/23] Update elbo test

---
 test/elbo.jl | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/test/elbo.jl b/test/elbo.jl
index fae4c2a1..f1ee971b 100644
--- a/test/elbo.jl
+++ b/test/elbo.jl
@@ -2,6 +2,7 @@ using Distributions
 using Pathfinder
 using Random
 using Test
+using Transducers
 
 @testset "ELBO estimation" begin
     @testset "elbo_and_samples" begin
@@ -30,7 +31,8 @@ using Test
         σs = [1e-3, 0.05, σ_target, 1.0, 1.1, 1.2, 5.0, 10.0]
         dists = Normal.(0, σs)
         rng = MersenneTwister(42)
-        lopt, elbo, ϕ, logqϕ = @inferred Pathfinder.maximize_elbo(rng, logp, dists, 100)
+        executor = SerialEx()
+        lopt, elbo, ϕ, logqϕ = @inferred Pathfinder.maximize_elbo(rng, logp, dists, 100, executor)
         @test lopt == 3
         @test elbo ≈ 0
     end

From aa0096981a58323039dd0ccc189bb01adfbcea1f Mon Sep 17 00:00:00 2001
From: Seth Axen <seth.axen@gmail.com>
Date: Mon, 25 Apr 2022 00:22:20 +0200
Subject: [PATCH 11/23] Fix name of type

---
 test/elbo.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/elbo.jl b/test/elbo.jl
index f1ee971b..0e9e0fa8 100644
--- a/test/elbo.jl
+++ b/test/elbo.jl
@@ -31,7 +31,7 @@ using Transducers
         σs = [1e-3, 0.05, σ_target, 1.0, 1.1, 1.2, 5.0, 10.0]
         dists = Normal.(0, σs)
         rng = MersenneTwister(42)
-        executor = SerialEx()
+        executor = Transducers.SequentialEx()
         lopt, elbo, ϕ, logqϕ = @inferred Pathfinder.maximize_elbo(rng, logp, dists, 100, executor)
         @test lopt == 3
         @test elbo ≈ 0

From 670276f43bb76ce9b1b5ec6e4b282876770157d8 Mon Sep 17 00:00:00 2001
From: Seth Axen <seth.axen@gmail.com>
Date: Tue, 26 Apr 2022 16:16:08 +0200
Subject: [PATCH 12/23] Don't nest multi-threading by default

---
 src/multipath.jl | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/src/multipath.jl b/src/multipath.jl
index 2c2f28fc..1e70a0a0 100644
--- a/src/multipath.jl
+++ b/src/multipath.jl
@@ -42,9 +42,10 @@ resulting draws better approximate draws from the target distribution ``p`` inst
 - `executor::Transducers.Executor`: Transducers.jl executor that determines if and how
     to run the single-path runs in parallel. If `rng` is known to be thread-safe, the
     default is `Transducers.PreferParallel(; basesize=1)` (parallel executation, defaulting
-    to multi-threading). Otherwise, it is `Transducers.SerialEx()` (no parallelization).
-- `executor_per_run::Transducers.Executor`: Transducers.jl executor used within each run to
-    parallelize PRNG calls. See [`pathfinder`](@ref) for a description.
+    to multi-threading). Otherwise, it is `Transducers.SequentialEx()` (no parallelization).
+- `executor_per_run::Transducers.Executor=Transducers.SequentialEx()`: Transducers.jl
+    executor used within each run to parallelize PRNG calls. Defaults to no parallelization.
+    See [`pathfinder`](@ref) for a description.
 - `kwargs...` : Remaining keywords are forwarded to [`pathfinder`](@ref).
 
 # Returns
@@ -90,7 +91,7 @@ function multipathfinder(
     ndraws_per_run::Int=5,
     rng::Random.AbstractRNG=Random.GLOBAL_RNG,
     executor::Transducers.Executor=_default_executor(rng; basesize=1),
-    executor_per_run=_default_executor(rng),
+    executor_per_run=Transducers.SequentialEx(),
     importance::Bool=true,
     kwargs...,
 )
@@ -103,7 +104,6 @@ function multipathfinder(
     logp(x) = -optim_fun.f(x, nothing)
 
     # run pathfinder independently from each starting point
-    # TODO: allow to be parallelized
     trans = Transducers.Map() do θ₀
         return pathfinder(
             optim_fun, θ₀, ndraws_per_run; rng, executor=executor_per_run, kwargs...

From 8bc0fce97c6b6011ccc40944fb9e7f26aff396ea Mon Sep 17 00:00:00 2001
From: Seth Axen <seth.axen@gmail.com>
Date: Tue, 26 Apr 2022 16:40:35 +0200
Subject: [PATCH 13/23] Work around type-instability

---
 src/elbo.jl | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

diff --git a/src/elbo.jl b/src/elbo.jl
index 9b86abb6..b39bb58c 100644
--- a/src/elbo.jl
+++ b/src/elbo.jl
@@ -1,15 +1,20 @@
 function maximize_elbo(rng, logp, dists, ndraws, executor)
-    iter = dists |> Transducers.Map() do dist
-        elbo_and_samples(rng, logp, dist, ndraws)
+    elbo_ϕ_logqϕ1 = elbo_and_samples(rng, logp, dists[begin], ndraws)
+    elbo_ϕ_logqϕ = similar(dists, typeof(elbo_ϕ_logqϕ1))
+    elbo_ϕ_logqϕ[begin] = elbo_ϕ_logqϕ1
+    iter = eachindex(dists, elbo_ϕ_logqϕ)[begin+1:end] |> Transducers.Map() do i
+        elbo_ϕ_logqϕ[i] = elbo_and_samples(rng, logp, dists[i], ndraws)
+        return nothing
     end
-    elbo_ϕ_logqϕ = Folds.collect(iter, executor)
+    Folds.collect(iter, executor)
     _, lopt = _findmax(elbo_ϕ_logqϕ |> Transducers.Map(first))
     return (lopt, elbo_ϕ_logqϕ[lopt]...)
 end
 
 function elbo_and_samples(rng, logp, dist, ndraws)
     ϕ, logqϕ = rand_and_logpdf(rng, dist, ndraws)
-    logpϕ = eachcol(ϕ) |> Transducers.Map(logp) |> Transducers.tcollect
+    logpϕ = similar(logqϕ)
+    logpϕ .= logp.(eachcol(ϕ))
     elbo = elbo_from_logpdfs(logpϕ, logqϕ)
     return elbo, ϕ, logqϕ
 end

From 5d3c0143ee6e1deac8d42e876a9678be938ed4e3 Mon Sep 17 00:00:00 2001
From: Seth Axen <seth.axen@gmail.com>
Date: Tue, 26 Apr 2022 16:41:39 +0200
Subject: [PATCH 14/23] Compute log-gradient in parallel

---
 src/optimize.jl   | 16 ++++++++--------
 src/singlepath.jl |  2 +-
 2 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/src/optimize.jl b/src/optimize.jl
index d901a077..34b17e26 100644
--- a/src/optimize.jl
+++ b/src/optimize.jl
@@ -29,16 +29,10 @@ function build_optim_problem(optim_fun, x₀; kwargs...)
     return GalacticOptim.OptimizationProblem(optim_fun, x₀, nothing; kwargs...)
 end
 
-function optimize_with_trace(prob, optimizer)
+function optimize_with_trace(prob, optimizer, executor=Transducers.SequentialEx())
     u0 = prob.u0
     fun = prob.f
     grad! = fun.grad
-    function ∇f(x)
-        ∇fx = similar(x)
-        grad!(∇fx, x, nothing)
-        rmul!(∇fx, -1)
-        return ∇fx
-    end
     # caches for the trace of x and f(x)
     xs = typeof(u0)[]
     fxs = typeof(fun.f(u0, nothing))[]
@@ -52,6 +46,12 @@ function optimize_with_trace(prob, optimizer)
     # NOTE: GalacticOptim doesn't have an interface for accessing the gradient trace,
     # so we need to recompute it ourselves
     # see https://github.com/SciML/GalacticOptim.jl/issues/149
-    ∇fxs = xs |> Transducers.Map(∇f) |> Transducers.tcollect
+    ∇fxs = [similar(u0) for _ in xs]
+    trans = Transducers.Map() do (∇fx, x)
+        grad!(∇fx, x, nothing)
+        rmul!(∇fx, -1)
+        return nothing
+    end
+    Folds.collect(zip(∇fxs, xs) |> trans, executor)
     return xs, fxs, ∇fxs
 end
diff --git a/src/singlepath.jl b/src/singlepath.jl
index a01b1e12..4b0b0921 100644
--- a/src/singlepath.jl
+++ b/src/singlepath.jl
@@ -109,7 +109,7 @@ function pathfinder(
     end
     logp(x) = -optim_prob.f.f(x, nothing)
     # compute trajectory
-    θs, logpθs, ∇logpθs = optimize_with_trace(optim_prob, optimizer)
+    θs, logpθs, ∇logpθs = optimize_with_trace(optim_prob, optimizer, executor)
     L = length(θs) - 1
     @assert L + 1 == length(logpθs) == length(∇logpθs)
 

From 81a8f38defce5d15e02a7ff532697a2c72997a97 Mon Sep 17 00:00:00 2001
From: Seth Axen <seth.axen@gmail.com>
Date: Tue, 26 Apr 2022 16:42:15 +0200
Subject: [PATCH 15/23] Run single-path sequentially by default

---
 src/singlepath.jl | 14 ++++++++------
 1 file changed, 8 insertions(+), 6 deletions(-)

diff --git a/src/singlepath.jl b/src/singlepath.jl
index 4b0b0921..c81cd41c 100644
--- a/src/singlepath.jl
+++ b/src/singlepath.jl
@@ -21,10 +21,12 @@ constructed using at most the previous `history_length` steps.
 # Keywords
 - `ad_backend=AD.ForwardDiffBackend()`: AbstractDifferentiation.jl AD backend.
 - `rng::Random.AbstractRNG`: The random number generator to be used for drawing samples
-- `executor::Transducers.Executor`: Transducers.jl executor that determines if and how
-    to perform ELBO computation in parallel. If `rng` is known to be thread-safe, the
-    default is `Transducers.PreferParallel()` (parallel executation, defaulting to
-    multi-threading). Otherwise, it is `Transducers.SerialEx()` (no parallelization).
+- `executor::Transducers.Executor=Transducers.SequentialEx()`: Transducers.jl executor that
+    determines if and how to perform ELBO computation in parallel. The default
+    (`SequentialEx()`) performs no parallelization. If `rng` is known to be thread-safe, and
+    the log-density function is known to have no internal state, then
+    `Transducers.PreferParallel()` may be used to parallelize log-density evaluation.
+    This is generally only faster for expensive log density functions.
 - `optimizer`: Optimizer to be used for constructing trajectory. Can be any optimizer
     compatible with GalacticOptim, so long as it supports callbacks. Defaults to
     `Optim.LBFGS(; m=$DEFAULT_HISTORY_LENGTH, linesearch=LineSearches.MoreThuente())`. See
@@ -70,7 +72,7 @@ function pathfinder(
     θ₀,
     ndraws;
     rng::Random.AbstractRNG=Random.GLOBAL_RNG,
-    executor::Transducers.Executor=_default_executor(rng),
+    executor::Transducers.Executor=Transducers.SequentialEx(),
     optimizer=DEFAULT_OPTIMIZER,
     history_length::Int=optimizer isa Optim.LBFGS ? optimizer.m : DEFAULT_HISTORY_LENGTH,
     ndraws_elbo::Int=5,
@@ -99,7 +101,7 @@ function pathfinder(
     optim_prob::GalacticOptim.OptimizationProblem,
     ndraws;
     rng::Random.AbstractRNG=Random.GLOBAL_RNG,
-    executor::Transducers.Executor=_default_executor(rng),
+    executor::Transducers.Executor=Transducers.SequentialEx(),
     optimizer=DEFAULT_OPTIMIZER,
     history_length::Int=optimizer isa Optim.LBFGS ? optimizer.m : DEFAULT_HISTORY_LENGTH,
     ndraws_elbo::Int=5,

From 345a735a96638f2551ae865622f03943f5a5e9ff Mon Sep 17 00:00:00 2001
From: Seth Axen <seth.axen@gmail.com>
Date: Tue, 26 Apr 2022 16:42:39 +0200
Subject: [PATCH 16/23] Update tests

---
 test/elbo.jl     | 22 +++++++++++++++++-----
 test/optimize.jl |  8 +++++++-
 2 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/test/elbo.jl b/test/elbo.jl
index 0e9e0fa8..c8635e22 100644
--- a/test/elbo.jl
+++ b/test/elbo.jl
@@ -30,10 +30,22 @@ using Transducers
         logp(x) = logpdf(target_dist, x[1])
         σs = [1e-3, 0.05, σ_target, 1.0, 1.1, 1.2, 5.0, 10.0]
         dists = Normal.(0, σs)
-        rng = MersenneTwister(42)
-        executor = Transducers.SequentialEx()
-        lopt, elbo, ϕ, logqϕ = @inferred Pathfinder.maximize_elbo(rng, logp, dists, 100, executor)
-        @test lopt == 3
-        @test elbo ≈ 0
+        if VERSION ≥ v"1.7.0"
+            executors = [SequentialEx(), ThreadedEx()]
+        else
+            executors = [SequentialEx()]
+        end
+        @testset "$executor" for executor in executors
+            rng = Random.seed!(42)
+            lopt, elbo, ϕ, logqϕ = @inferred Pathfinder.maximize_elbo(rng, logp, dists, 100, executor)
+            @test lopt == 3
+            @test elbo ≈ 0
+            rng = Random.seed!(42)
+            lopt2, elbo2, ϕ2, logqϕ2 = Pathfinder.maximize_elbo(rng, logp, dists, 100, executor)
+            @test lopt2 == lopt
+            @test elbo2 == elbo
+            @test ϕ2 ≈ ϕ
+            @test logqϕ ≈ logqϕ2
+        end
     end
 end
diff --git a/test/optimize.jl b/test/optimize.jl
index f3bac899..999602aa 100644
--- a/test/optimize.jl
+++ b/test/optimize.jl
@@ -6,6 +6,7 @@ using NLopt
 using Optim
 using Pathfinder
 using Test
+using Transducers
 
 include("test_utils.jl")
 
@@ -65,7 +66,7 @@ end
         Optim.BFGS(), Optim.LBFGS(), Optim.ConjugateGradient(), NLopt.Opt(:LD_LBFGS, n)
     ]
     @testset "$(typeof(optimizer))" for optimizer in optimizers
-        xs, fxs, ∇fxs = Pathfinder.optimize_with_trace(prob, optimizer)
+        xs, fxs, ∇fxs = Pathfinder.optimize_with_trace(prob, optimizer, SequentialEx())
         @test xs[1] ≈ x0
         @test xs[end] ≈ μ
         @test fxs ≈ f.(xs)
@@ -80,5 +81,10 @@ end
             @test Optim.x_trace(res) ≈ xs
             @test Optim.minimizer(res) ≈ xs[end]
         end
+
+        xs2, fxs2, ∇fxs2 = Pathfinder.optimize_with_trace(prob, optimizer, ThreadedEx())
+        @test xs2 ≈ xs
+        @test fxs2 ≈ fxs
+        @test ∇fxs2 ≈ ∇fxs
     end
 end

From e1a5329751f443250be20de0631a26c9632b29cd Mon Sep 17 00:00:00 2001
From: Seth Axen <seth.axen@gmail.com>
Date: Tue, 26 Apr 2022 16:44:13 +0200
Subject: [PATCH 17/23] Run formatter

---
 src/elbo.jl      | 9 +++++----
 test/elbo.jl     | 8 ++++++--
 test/woodbury.jl | 3 ++-
 3 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/src/elbo.jl b/src/elbo.jl
index b39bb58c..9490657b 100644
--- a/src/elbo.jl
+++ b/src/elbo.jl
@@ -2,10 +2,11 @@ function maximize_elbo(rng, logp, dists, ndraws, executor)
     elbo_ϕ_logqϕ1 = elbo_and_samples(rng, logp, dists[begin], ndraws)
     elbo_ϕ_logqϕ = similar(dists, typeof(elbo_ϕ_logqϕ1))
     elbo_ϕ_logqϕ[begin] = elbo_ϕ_logqϕ1
-    iter = eachindex(dists, elbo_ϕ_logqϕ)[begin+1:end] |> Transducers.Map() do i
-        elbo_ϕ_logqϕ[i] = elbo_and_samples(rng, logp, dists[i], ndraws)
-        return nothing
-    end
+    iter =
+        eachindex(dists, elbo_ϕ_logqϕ)[(begin + 1):end] |> Transducers.Map() do i
+            elbo_ϕ_logqϕ[i] = elbo_and_samples(rng, logp, dists[i], ndraws)
+            return nothing
+        end
     Folds.collect(iter, executor)
     _, lopt = _findmax(elbo_ϕ_logqϕ |> Transducers.Map(first))
     return (lopt, elbo_ϕ_logqϕ[lopt]...)
diff --git a/test/elbo.jl b/test/elbo.jl
index c8635e22..2d533907 100644
--- a/test/elbo.jl
+++ b/test/elbo.jl
@@ -37,11 +37,15 @@ using Transducers
         end
         @testset "$executor" for executor in executors
             rng = Random.seed!(42)
-            lopt, elbo, ϕ, logqϕ = @inferred Pathfinder.maximize_elbo(rng, logp, dists, 100, executor)
+            lopt, elbo, ϕ, logqϕ = @inferred Pathfinder.maximize_elbo(
+                rng, logp, dists, 100, executor
+            )
             @test lopt == 3
             @test elbo ≈ 0
             rng = Random.seed!(42)
-            lopt2, elbo2, ϕ2, logqϕ2 = Pathfinder.maximize_elbo(rng, logp, dists, 100, executor)
+            lopt2, elbo2, ϕ2, logqϕ2 = Pathfinder.maximize_elbo(
+                rng, logp, dists, 100, executor
+            )
             @test lopt2 == lopt
             @test elbo2 == elbo
             @test ϕ2 ≈ ϕ
diff --git a/test/woodbury.jl b/test/woodbury.jl
index b8b56fad..c4a9461d 100644
--- a/test/woodbury.jl
+++ b/test/woodbury.jl
@@ -160,7 +160,8 @@ end
             @test @inferred(quad(W, X)) ≈ quad(PDMats.PDMat(Symmetric(Wmat)), X)
 
             U = randn(T, n, 100)
-            @test quad(W, Pathfinder.invunwhiten!(similar(U), W, U)) ≈ vec(sum(abs2, U; dims=1))
+            @test quad(W, Pathfinder.invunwhiten!(similar(U), W, U)) ≈
+                vec(sum(abs2, U; dims=1))
         end
 
         @testset "PDMats.invquad" begin

From a71a6b14f6584c69a37605bd67cf25b83b184f8e Mon Sep 17 00:00:00 2001
From: Seth Axen <seth.axen@gmail.com>
Date: Tue, 26 Apr 2022 22:06:31 +0200
Subject: [PATCH 18/23] Test with more rngs and executors

---
 test/multipath.jl  | 77 +++++++++++++++++++++++++++-------------------
 test/singlepath.jl | 39 ++++++++++++++++++-----
 2 files changed, 77 insertions(+), 39 deletions(-)

diff --git a/test/multipath.jl b/test/multipath.jl
index 885287af..e76b4b3d 100644
--- a/test/multipath.jl
+++ b/test/multipath.jl
@@ -14,45 +14,58 @@ using Test
         nruns = 20
         ndraws = 1000_000
         ndraws_per_run = ndraws ÷ nruns
+        ndraws_elbo = 100
         Σ = rand_pd_mat(rng, Float64, n)
         μ = randn(rng, n)
         d = MvNormal(μ, Σ)
         logp(x) = logpdf(d, x)
         ∇logp(x) = ForwardDiff.gradient(logp, x)
         x₀s = [rand(rng, Uniform(-2, 2), n) for _ in 1:nruns]
-        rng = MersenneTwister(76)
-        q, ϕ, component_ids = multipathfinder(
-            logp, ∇logp, x₀s, ndraws; ndraws_elbo=100, ndraws_per_run, rng
-        )
-        @test q isa MixtureModel
-        @test ncomponents(q) == nruns
-        @test Distributions.component_type(q) <: MvNormal
-        @test ϕ isa AbstractMatrix
-        @test size(ϕ) == (n, ndraws)
-        @test component_ids isa Vector{Int}
-        @test length(component_ids) == ndraws
-        @test extrema(component_ids) == (1, nruns)
-        μ_hat = mean(ϕ; dims=2)
-        Σ_hat = cov(ϕ .- μ_hat; dims=2, corrected=false)
-        # adapted from the MvNormal tests
-        # allow for 15x disagreement in atol, since this method is approximate
-        multiplier = 15
-        for i in 1:n
-            atol = sqrt(Σ[i, i] / ndraws) * 8 * multiplier
-            @test μ_hat[i] ≈ μ[i] atol = atol
-        end
-        for i in 1:n, j in 1:n
-            atol = sqrt(Σ[i, i] * Σ[j, j] / ndraws) * 10 * multiplier
-            @test isapprox(Σ_hat[i, j], Σ[i, j], atol=atol)
-        end
+        @testset for rng in [MersenneTwister(), Random.default_rng()]
+            executor = rng isa MersenneTwister ? SequentialEx() : ThreadedEx()
+
+            Random.seed!(rng, 76)
+            q, ϕ, component_ids = multipathfinder(
+                logp, ∇logp, x₀s, ndraws; ndraws_elbo, ndraws_per_run, rng, executor
+            )
+            @test q isa MixtureModel
+            @test ncomponents(q) == nruns
+            @test Distributions.component_type(q) <: MvNormal
+            @test ϕ isa AbstractMatrix
+            @test size(ϕ) == (n, ndraws)
+            @test component_ids isa Vector{Int}
+            @test length(component_ids) == ndraws
+            @test extrema(component_ids) == (1, nruns)
+            μ_hat = mean(ϕ; dims=2)
+            Σ_hat = cov(ϕ .- μ_hat; dims=2, corrected=false)
+            # adapted from the MvNormal tests
+            # allow for 15x disagreement in atol, since this method is approximate
+            multiplier = 15
+            for i in 1:n
+                atol = sqrt(Σ[i, i] / ndraws) * 8 * multiplier
+                @test μ_hat[i] ≈ μ[i] atol = atol
+            end
+            for i in 1:n, j in 1:n
+                atol = sqrt(Σ[i, i] * Σ[j, j] / ndraws) * 10 * multiplier
+                @test isapprox(Σ_hat[i, j], Σ[i, j], atol=atol)
+            end
+
+            Random.seed!(rng, 76)
+            q2, ϕ2, component_ids2 = multipathfinder(
+                logp, x₀s, ndraws; ndraws_elbo, ndraws_per_run, rng, executor
+            )
+            @test q2 == q
+            @test ϕ2 == ϕ
+            @test component_ids2 == component_ids
 
-        rng = MersenneTwister(76)
-        ad_backend = AD.ReverseDiffBackend()
-        q2, ϕ2, component_ids2 = multipathfinder(
-            logp, x₀s, ndraws; ndraws_elbo=100, ndraws_per_run, rng, ad_backend
-        )
-        for (c1, c2) in zip(q.components, q2.components)
-            @test c1 ≈ c2 atol = 1e-6
+            Random.seed!(rng, 76)
+            ad_backend = AD.ReverseDiffBackend()
+            q3, ϕ3, component_ids3 = multipathfinder(
+                logp, x₀s, ndraws; ndraws_elbo, ndraws_per_run, rng, executor, ad_backend
+            )
+            for (c1, c2) in zip(q.components, q3.components)
+                @test c1 ≈ c2 atol = 1e-6
+            end
         end
     end
     @testset "errors if no gradient provided" begin
diff --git a/test/singlepath.jl b/test/singlepath.jl
index 6fda9f90..3a9dd7fd 100644
--- a/test/singlepath.jl
+++ b/test/singlepath.jl
@@ -4,8 +4,10 @@ using ForwardDiff
 using GalacticOptim
 using LinearAlgebra
 using Pathfinder
+using Random
 using ReverseDiff
 using Test
+using Transducers
 
 @testset "single path pathfinder" begin
     @testset "IsoNormal" begin
@@ -13,9 +15,12 @@ using Test
         logp(x) = -sum(abs2, x) / 2
         ∇logp(x) = -x
         ndraws = 100
-        @testset for n in [1, 5, 10, 100]
+        @testset for n in [1, 5, 10, 100], rng in [MersenneTwister(), Random.default_rng()]
+            executor = rng isa MersenneTwister ? SequentialEx() : ThreadedEx()
+
             x0 = randn(n)
-            q, ϕ, logqϕ = @inferred pathfinder(logp, ∇logp, x0, ndraws)
+            Random.seed!(rng, 42)
+            q, ϕ, logqϕ = @inferred pathfinder(logp, ∇logp, x0, ndraws; rng, executor)
             @test q isa MvNormal
             @test q.μ ≈ zeros(n)
             @test q.Σ isa Pathfinder.WoodburyPDMat
@@ -25,8 +30,14 @@ using Test
             @test size(ϕ) == (n, ndraws)
             @test logqϕ ≈ logpdf(q, ϕ)
 
-            q2, ϕ2, logqϕ2 = pathfinder(logp, ∇logp, x0, 2)
-            @test size(ϕ2) == (n, 2)
+            Random.seed!(rng, 42)
+            q2, ϕ2, logqϕ2 = pathfinder(logp, ∇logp, x0, ndraws; rng, executor)
+            @test q2 == q
+            @test ϕ2 == ϕ
+            @test logqϕ2 == logqϕ
+
+            q3, ϕ3, logqϕ3 = pathfinder(logp, ∇logp, x0, 2; executor)
+            @test size(ϕ3) == (n, 2)
         end
     end
     @testset "MvNormal" begin
@@ -43,10 +54,24 @@ using Test
         logp(x) = -dot(x, P, x) / 2
         ∇logp(x) = -(P * x)
         x₀ = [2.08, 3.77, -1.26, -0.97, -3.91]
-        rng = MersenneTwister(38)
         ad_backend = AD.ReverseDiffBackend()
-        q, _, _ = @inferred pathfinder(logp, x₀, 10; rng, ndraws_elbo=100, ad_backend)
-        @test q.Σ ≈ Σ rtol = 1e-1
+        ndraws_elbo = 100
+        @testset for rng in [MersenneTwister(), Random.default_rng()]
+            executor = rng isa MersenneTwister ? SequentialEx() : ThreadedEx()
+
+            Random.seed!(rng, 38)
+            q, ϕ, logqϕ = @inferred pathfinder(
+                logp, x₀, 10; rng, ndraws_elbo, ad_backend, executor
+            )
+            @test q.Σ ≈ Σ rtol = 1e-1
+            Random.seed!(rng, 38)
+            q2, ϕ2, logqϕ2 = pathfinder(
+                logp, x₀, 10; rng, ndraws_elbo, ad_backend, executor
+            )
+            @test q2 == q
+            @test ϕ2 == ϕ
+            @test logqϕ2 == logqϕ
+        end
     end
     @testset "errors if no gradient provided" begin
         logp(x) = -sum(abs2, x) / 2

From 26beeb711988378c8c70fc0c04d873b88edacda8 Mon Sep 17 00:00:00 2001
From: Seth Axen <seth.axen@gmail.com>
Date: Tue, 26 Apr 2022 22:10:11 +0200
Subject: [PATCH 19/23] Avoid double-testing pre-1.7

---
 test/multipath.jl  |  7 ++++++-
 test/singlepath.jl | 14 ++++++++++++--
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/test/multipath.jl b/test/multipath.jl
index e76b4b3d..6e9f398d 100644
--- a/test/multipath.jl
+++ b/test/multipath.jl
@@ -21,7 +21,12 @@ using Test
         logp(x) = logpdf(d, x)
         ∇logp(x) = ForwardDiff.gradient(logp, x)
         x₀s = [rand(rng, Uniform(-2, 2), n) for _ in 1:nruns]
-        @testset for rng in [MersenneTwister(), Random.default_rng()]
+        rngs = if VERSION ≥ v"1.7"
+            [MersenneTwister(), Random.default_rng()]
+        else
+            [MersenneTwister()]
+        end
+        @testset for rng in rngs
             executor = rng isa MersenneTwister ? SequentialEx() : ThreadedEx()
 
             Random.seed!(rng, 76)
diff --git a/test/singlepath.jl b/test/singlepath.jl
index 3a9dd7fd..f8c8e4c9 100644
--- a/test/singlepath.jl
+++ b/test/singlepath.jl
@@ -15,7 +15,12 @@ using Transducers
         logp(x) = -sum(abs2, x) / 2
         ∇logp(x) = -x
         ndraws = 100
-        @testset for n in [1, 5, 10, 100], rng in [MersenneTwister(), Random.default_rng()]
+        rngs = if VERSION ≥ v"1.7"
+            [MersenneTwister(), Random.default_rng()]
+        else
+            [MersenneTwister()]
+        end
+        @testset for n in [1, 5, 10, 100], rng in rngs
             executor = rng isa MersenneTwister ? SequentialEx() : ThreadedEx()
 
             x0 = randn(n)
@@ -56,7 +61,12 @@ using Transducers
         x₀ = [2.08, 3.77, -1.26, -0.97, -3.91]
         ad_backend = AD.ReverseDiffBackend()
         ndraws_elbo = 100
-        @testset for rng in [MersenneTwister(), Random.default_rng()]
+        rngs = if VERSION ≥ v"1.7"
+            [MersenneTwister(), Random.default_rng()]
+        else
+            [MersenneTwister()]
+        end
+        @testset for rng in rngs
             executor = rng isa MersenneTwister ? SequentialEx() : ThreadedEx()
 
             Random.seed!(rng, 38)

From ec84985183f83afa2d81d6149d804262ad98a88a Mon Sep 17 00:00:00 2001
From: Seth Axen <seth.axen@gmail.com>
Date: Tue, 26 Apr 2022 22:13:10 +0200
Subject: [PATCH 20/23] Test with multiple threads

---
 .github/workflows/CI.yml | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
index d9b0da57..ac431507 100644
--- a/.github/workflows/CI.yml
+++ b/.github/workflows/CI.yml
@@ -14,7 +14,7 @@ concurrency:
   cancel-in-progress: ${{ startsWith(github.ref, 'refs/pull/') }}
 jobs:
   test:
-    name: Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }} - ${{ github.event_name }}
+    name: Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }} - ${{ matrix.num_threads }} threads - ${{ github.event_name }}
     runs-on: ${{ matrix.os }}
     strategy:
       fail-fast: false
@@ -26,6 +26,9 @@ jobs:
           - ubuntu-latest
         arch:
           - x64
+        num_threads:
+          - 1
+          - 2
     steps:
       - uses: actions/checkout@v2
       - uses: julia-actions/setup-julia@v1
@@ -35,6 +38,8 @@ jobs:
       - uses: julia-actions/cache@v1
       - uses: julia-actions/julia-buildpkg@v1
       - uses: julia-actions/julia-runtest@v1
+        env:
+          JULIA_NUM_THREADS: ${{ matrix.num_threads }}
       - uses: julia-actions/julia-processcoverage@v1
       - uses: codecov/codecov-action@v2
         with:

From 251fd59056eb186c241144efe13b7acfec8786da Mon Sep 17 00:00:00 2001
From: Seth Axen <seth.axen@gmail.com>
Date: Tue, 26 Apr 2022 22:20:22 +0200
Subject: [PATCH 21/23] Add transducer-specific tests

---
 test/runtests.jl    |  1 +
 test/transducers.jl | 29 +++++++++++++++++++++++++++++
 2 files changed, 30 insertions(+)
 create mode 100644 test/transducers.jl

diff --git a/test/runtests.jl b/test/runtests.jl
index 309d8312..022e828a 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -2,6 +2,7 @@ using Pathfinder
 using Test
 
 @testset "Pathfinder.jl" begin
+    include("transducers.jl")
     include("woodbury.jl")
     include("optimize.jl")
     include("inverse_hessian.jl")
diff --git a/test/transducers.jl b/test/transducers.jl
new file mode 100644
index 00000000..1d7bb935
--- /dev/null
+++ b/test/transducers.jl
@@ -0,0 +1,29 @@
+using Pathfinder
+using Random
+using Test
+using Transducers
+
+@testset "transducers integration" begin
+    @testset "_default_executor" begin
+        rng = MersenneTwister(42)
+        @test Pathfinder._default_executor(MersenneTwister(42); basesize=2) ===
+            SequentialEx()
+        if VERSION ≥ v"1.7.0"
+            @test Pathfinder._default_executor(Random.GLOBAL_RNG; basesize=1) ===
+                PreferParallel(; basesize=1)
+            @test Pathfinder._default_executor(Random.default_rng(); basesize=1) ===
+                PreferParallel(; basesize=1)
+        else
+            @test Pathfinder._default_executor(Random.GLOBAL_RNG; basesize=1) ===
+                SequentialEx()
+            @test Pathfinder._default_executor(Random.default_rng(); basesize=1) ===
+                SequentialEx()
+        end
+    end
+
+    @testset "_findmax" begin
+        x = randn(100)
+        @test Pathfinder._findmax(x) == findmax(x)
+        @test Pathfinder._findmax(x |> Map(sin)) == findmax(sin.(x))
+    end
+end

From 518c3c4018ef99501d9a37c0f932fe76d6fce9f4 Mon Sep 17 00:00:00 2001
From: Seth Axen <seth.axen@gmail.com>
Date: Tue, 26 Apr 2022 23:00:56 +0200
Subject: [PATCH 22/23] Use reducer instead of collector

---
 src/elbo.jl     | 11 +++++------
 src/optimize.jl |  2 +-
 2 files changed, 6 insertions(+), 7 deletions(-)

diff --git a/src/elbo.jl b/src/elbo.jl
index 9490657b..87e26a5b 100644
--- a/src/elbo.jl
+++ b/src/elbo.jl
@@ -2,12 +2,11 @@ function maximize_elbo(rng, logp, dists, ndraws, executor)
     elbo_ϕ_logqϕ1 = elbo_and_samples(rng, logp, dists[begin], ndraws)
     elbo_ϕ_logqϕ = similar(dists, typeof(elbo_ϕ_logqϕ1))
     elbo_ϕ_logqϕ[begin] = elbo_ϕ_logqϕ1
-    iter =
-        eachindex(dists, elbo_ϕ_logqϕ)[(begin + 1):end] |> Transducers.Map() do i
-            elbo_ϕ_logqϕ[i] = elbo_and_samples(rng, logp, dists[i], ndraws)
-            return nothing
-        end
-    Folds.collect(iter, executor)
+    @views Folds.map!(
+        elbo_ϕ_logqϕ[(begin + 1):end], dists[(begin + 1):end], executor
+    ) do dist
+        return elbo_and_samples(rng, logp, dist, ndraws)
+    end
     _, lopt = _findmax(elbo_ϕ_logqϕ |> Transducers.Map(first))
     return (lopt, elbo_ϕ_logqϕ[lopt]...)
 end
diff --git a/src/optimize.jl b/src/optimize.jl
index 34b17e26..2621e969 100644
--- a/src/optimize.jl
+++ b/src/optimize.jl
@@ -52,6 +52,6 @@ function optimize_with_trace(prob, optimizer, executor=Transducers.SequentialEx(
         rmul!(∇fx, -1)
         return nothing
     end
-    Folds.collect(zip(∇fxs, xs) |> trans, executor)
+    Transducers.transduce(trans, (_...,) -> nothing, nothing, zip(∇fxs, xs), executor)
     return xs, fxs, ∇fxs
 end

From 27e542a2f048dd9050e48134ac58d7405991d12a Mon Sep 17 00:00:00 2001
From: Seth Axen <seth.axen@gmail.com>
Date: Tue, 26 Apr 2022 23:52:30 +0200
Subject: [PATCH 23/23] Increment patch number

---
 Project.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Project.toml b/Project.toml
index 00a80f51..7a002e5e 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "Pathfinder"
 uuid = "b1d3bc72-d0e7-4279-b92f-7fa5d6d2d454"
 authors = ["Seth Axen <seth.axen@gmail.com> and contributors"]
-version = "0.3.5"
+version = "0.3.6"
 
 [deps]
 AbstractDifferentiation = "c29ec348-61ec-40c8-8164-b8c60e9d9f3d"