From 04f615350fad79c924aaa73a9db7b1d5bf780657 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Mon, 25 May 2026 11:46:22 +0200 Subject: [PATCH 1/7] refactor: replace Transducers/Folds with OhMyThreads Drop the unmaintained Transducers.jl and Folds.jl dependencies and route parallel iteration through OhMyThreads.jl instead. The public-facing `executor::Transducers.Executor` keyword of `pathfinder` and `multipathfinder` is replaced by `ntasks::Int=1` (and `ntasks_per_run` on `multipathfinder`); per-element RNG seeds are pre-drawn from the user's `rng` so output is reproducible regardless of `ntasks`. Adds explicit reproducibility-across-`ntasks` tests for both explicit and default RNGs. Co-Authored-By: Claude Opus 4.7 --- Project.toml | 8 +- docs/inventories/Transducers.toml | 11 -- docs/make.jl | 6 +- docs/src/index.md | 2 +- src/Pathfinder.jl | 5 +- src/elbo.jl | 44 ++++++-- src/multipath.jl | 172 +++++++++++++++++++++--------- src/mvnormal.jl | 4 +- src/singlepath.jl | 16 ++- src/transducers.jl | 18 ---- src/utils.jl | 24 +++++ test/Project.toml | 4 +- test/elbo.jl | 14 +-- test/multipath.jl | 50 +++++++-- test/runtests.jl | 2 +- test/singlepath.jl | 55 +++++++--- test/transducers.jl | 15 --- test/utils.jl | 14 +++ 18 files changed, 302 insertions(+), 162 deletions(-) delete mode 100644 docs/inventories/Transducers.toml delete mode 100644 src/transducers.jl create mode 100644 src/utils.jl delete mode 100644 test/transducers.jl create mode 100644 test/utils.jl diff --git a/Project.toml b/Project.toml index 56d130307..5ac6670cb 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "Pathfinder" uuid = "b1d3bc72-d0e7-4279-b92f-7fa5d6d2d454" -version = "0.9.31" +version = "0.10.0" authors = ["Seth Axen and contributors"] [workspace] @@ -9,11 +9,11 @@ projects = ["test"] [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" -Folds = "41a02a25-b8f0-4f67-bc48-60067656b558" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" IrrationalConstants = "92d709cd-6900-40b7-9082-c6be49f344b6" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c" +OhMyThreads = "67456a42-1dca-4109-a031-0a68de7e3ad5" Optim = "429524aa-4258-5aef-a3af-852621145aeb" Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba" OptimizationBase = "bca83a33-5cc9-4baa-983d-23429ab6bcbb" @@ -26,7 +26,6 @@ Requires = "ae029012-a4dd-5104-9daa-d747884805df" SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" -Transducers = "28d57a85-8fef-5791-bfe6-a80928e7c999" [weakdeps] AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001" @@ -47,12 +46,12 @@ Accessors = "0.1.12" Distributions = "0.25.87" DynamicHMC = "3.4.0" DynamicPPL = "0.40.6" -Folds = "0.2.9" ForwardDiff = "0.10.26, 1" IrrationalConstants = "0.1.1, 0.2" LinearAlgebra = "1" LogDensityProblems = "2.1.0" MCMCChains = "6.0.2, 7" +OhMyThreads = "0.8.5" Optim = "1.7.2, 2.0" Optimization = "3.22, 4, 5" OptimizationBase = "0.0.5, 1, 2, 3, 4, 5.0" @@ -65,6 +64,5 @@ Requires = "1.1.3" SciMLBase = "2.30" Statistics = "1" StatsBase = "0.33.17, 0.34" -Transducers = "0.4.81" Turing = "0.43" julia = "1.10" diff --git a/docs/inventories/Transducers.toml b/docs/inventories/Transducers.toml deleted file mode 100644 index 0411804b6..000000000 --- a/docs/inventories/Transducers.toml +++ /dev/null @@ -1,11 +0,0 @@ -# DocInventory version 1 -project = "Transducers.jl" -version = "0.4.84" - -# Filtered to just the types we link to -[[jl.type]] -name = "Transducers.PreferParallel" -uri = "reference/manual/#$" -[[jl.type]] -name = "Transducers.SequentialEx" -uri = "reference/manual/#$" diff --git a/docs/make.jl b/docs/make.jl index c00865816..9993e48de 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -41,12 +41,8 @@ links = InterLinks( ), "Optim" => "https://julianlsolvers.github.io/Optim.jl/stable/", "Optimization" => "https://docs.sciml.ai/Optimization/stable/", + "OhMyThreads" => "https://juliafolds2.github.io/OhMyThreads.jl/stable/", "PSIS" => "https://julia.arviz.org/PSIS/stable/", - "Transducers" => ( - "https://juliafolds2.github.io/Transducers.jl/stable/", # not built for a while - "https://juliafolds2.github.io/Transducers.jl/dev/objects.inv", - joinpath(@__DIR__, "inventories", "Transducers.toml"), - ), ) makedocs(; diff --git a/docs/src/index.md b/docs/src/index.md index 2f51db662..855cb89de 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -41,7 +41,7 @@ Pathfinder uses several packages for extended functionality: - [Optimization.jl](https://optimization.sciml.ai/stable/): This allows the L-BFGS optimizer to be replaced with any of the many Optimization-compatible optimizers and supports use of callbacks. Note that any changes made to Pathfinder using these features would be experimental. - [ADTypes.jl](https://sciml.github.io/ADTypes.jl/stable/): Supports specifying the automatic differentiation engine to be used for computing gradient and Hessian, if needed. -- [Transducers.jl](https://juliafolds.github.io/Transducers.jl/stable/): parallelization support +- [OhMyThreads.jl](https://juliafolds2.github.io/OhMyThreads.jl/stable/): parallelization support - [Distributions.jl](https://juliastats.org/Distributions.jl/stable/)/[PDMats.jl](https://github.com/JuliaStats/PDMats.jl): fits can be used anywhere a `Distribution` can be used - [LogDensityProblems.jl](https://www.tamaspapp.eu/LogDensityProblems.jl/stable/): defining the log-density function, gradient, and Hessian - [ProgressLogging.jl](https://julialogging.github.io/ProgressLogging.jl/stable/): In Pluto, Juno, and VSCode, nested progress bars are shown. In the REPL, use TerminalLoggers.jl to get progress bars. diff --git a/src/Pathfinder.jl b/src/Pathfinder.jl index e1ab7f8a3..f71d1be67 100644 --- a/src/Pathfinder.jl +++ b/src/Pathfinder.jl @@ -2,10 +2,10 @@ module Pathfinder using ADTypes: ADTypes using Distributions: Distributions -using Folds: Folds using IrrationalConstants: log2π using LinearAlgebra using LogDensityProblems: LogDensityProblems +using OhMyThreads: OhMyThreads using Optim: Optim, LineSearches using Optimization: Optimization using OptimizationOptimJL: OptimizationOptimJL @@ -17,7 +17,6 @@ using Requires: Requires using SciMLBase: SciMLBase using Statistics: Statistics using StatsBase: StatsBase -using Transducers: Transducers if isdefined(Optimization, :OptimizationState) using Optimization: OptimizationState @@ -45,7 +44,7 @@ end # We depend on Optim, and Optim depends on ForwardDiff, so we can offer it as a default. default_ad() = ADTypes.AutoForwardDiff() -include("transducers.jl") +include("utils.jl") include("woodbury.jl") include("optimize.jl") include("inverse_hessian.jl") diff --git a/src/elbo.jl b/src/elbo.jl index e7190e57b..727496308 100644 --- a/src/elbo.jl +++ b/src/elbo.jl @@ -1,13 +1,37 @@ -function maximize_elbo(rng, logp, dists, ndraws, executor) - EE = Core.Compiler.return_type( - elbo_and_samples, Tuple{typeof(rng),typeof(logp),eltype(dists),Int} - ) - estimates = similar(dists, EE) - isempty(estimates) && return 0, estimates - Folds.map!(estimates, dists, executor) do dist - return elbo_and_samples(rng, logp, dist, ndraws) +function maximize_elbo(rng, logp, dists, ndraws, ntasks::Int) + seeds = rand!(rng, similar(dists, UInt64)) + copy_rng = copy(rng) + nchunks = min(length(dists), ntasks) + estimates = if nchunks > 1 && Threads.nthreads() > 1 + nrngs = min(nchunks, Threads.nthreads()) + rng_pool = Channel{typeof(copy_rng)}(nrngs) + put!(rng_pool, copy_rng) + for _ in 2:nrngs + put!(rng_pool, copy(rng)) + end + OhMyThreads.tmapreduce( + vcat, + OhMyThreads.chunks(eachindex(dists, seeds); n=nchunks); + chunking=false, + ) do chunk + chunk_rng = take!(rng_pool) + try + return map(chunk) do i + Random.seed!(chunk_rng, seeds[i]) + return elbo_and_samples(chunk_rng, logp, dists[i], ndraws) + end + finally + put!(rng_pool, chunk_rng) + end + end + else + map(dists, seeds) do dist, seed + Random.seed!(copy_rng, seed) + return elbo_and_samples(copy_rng, logp, dist, ndraws) + end end - _, iteration_opt = _findmax(estimates |> Transducers.Map(est -> est.value)) + isempty(estimates) && return 0, estimates + _, iteration_opt = _findmax_skipnan(est -> est.value, estimates) return iteration_opt, estimates end @@ -17,7 +41,7 @@ function elbo_and_samples(rng, logp, dist, ndraws) logpϕ .= logp.(eachcol(ϕ)) logr = logpϕ - logqϕ elbo = Statistics.mean(logr) - elbo_se = sqrt(Statistics.var(logr) / length(logr)) + elbo_se = sqrt(Statistics.var(logr; mean = elbo) / length(logr)) return ELBOEstimate(elbo, elbo_se, ϕ, logpϕ, logqϕ, logr) end diff --git a/src/multipath.jl b/src/multipath.jl index ed5dd1f87..0d3df3107 100644 --- a/src/multipath.jl +++ b/src/multipath.jl @@ -101,15 +101,13 @@ $(_ARGUMENT_DOCSTRING) - `importance::Bool=true`: Perform Pareto smoothed importance resampling of draws. - `rng::AbstractRNG=Random.default_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, defaulting to - [`Transducers.SequentialEx()`](@extref `Transducers.SequentialEx`). If a transducer for - multi-threaded computation is selected, you must first verify that `rng` and the log - density function are thread-safe. -- `executor_per_run::Transducers.Executor`: Transducers.jl executor used within each run to - parallelize PRNG calls, defaulting to - [`Transducers.SequentialEx()`](@extref `Transducers.SequentialEx`). See - [`pathfinder`](@ref) for further description. +- `ntasks::Int=1`: maximum number of parallel tasks used to run the single-path + Pathfinder runs and to evaluate the target log density across draws. The default + `ntasks = 1` runs sequentially; larger values parallelize across runs and across + importance-resampling log-density evaluations, in which case the log-density function + must be thread-safe. Results are reproducible regardless of `ntasks`. +- `ntasks_per_run::Int=1`: same as `ntasks`, but applied within each single-path run for + parallelizing the ELBO evaluation. See [`pathfinder`](@ref) for details. - `kwargs...` : Remaining keywords are forwarded to [`pathfinder`](@ref). # Returns @@ -141,64 +139,142 @@ function multipathfinder( rng::Random.AbstractRNG=Random.default_rng(), history_length::Int=DEFAULT_HISTORY_LENGTH, optimizer=default_optimizer(history_length), - executor::Transducers.Executor=Transducers.SequentialEx(), - executor_per_run=Transducers.SequentialEx(), + ntasks::Int=1, + ntasks_per_run::Int=1, importance::Bool=true, kwargs..., ) - if init === nothing + _init = if init === nothing nruns > 0 || throw( ArgumentError("A positive `nruns` must be set or `init` must be provided.") ) - _init = fill(init, nruns) + fill(init, nruns) else - _init = init + init end + nruns = length(_init) if ndraws > ndraws_per_run * nruns @warn "More draws requested than total number of draws across replicas. Draws will not be unique." end logp(x) = -optim_fun.f(x, nothing) # run pathfinder independently from each starting point - trans = Transducers.Map() do (init_i, optimizer_i) - return pathfinder( - optim_fun; - rng, - history_length, - optimizer=optimizer_i, - ndraws=ndraws_per_run, - init=init_i, - executor=executor_per_run, - ndraws_elbo, - kwargs..., - ) + run_seeds = rand!(rng, similar(_init, UInt64)) + copy_rng = copy(rng) + nchunks = min(nruns, ntasks) + threaded = nchunks > 1 && Threads.nthreads() > 1 + + pathfinder_results = ProgressLogging.@withprogress name = "Multi-path Pathfinder" begin + progress_taskref = Ref{Task}() + progress_channel = Channel{Bool}( + min(nruns, 1_000); spawn=true, taskref=progress_taskref + ) do ch + # Throttle progress logs as generally the logging system is not super performant: + # - At most 1 progress log every 0.1 seconds + # - At most 1 progress log every 0.5% progress + progress_step = max(1, cld(nruns, 200)) + next_logged = progress_step + next_time = time() + 0.1 + + completed = 0 + while take!(ch) + completed += 1 + now = time() + if completed >= next_logged && now >= next_time + ProgressLogging.@logprogress completed / nruns + next_logged = completed + progress_step + next_time = now + 0.1 + end + end + end + + try + if threaded + rng_pool = Channel{typeof(copy_rng)}(nchunks) + put!(rng_pool, copy_rng) + for _ in 2:nchunks + put!(rng_pool, copy(rng)) + end + OhMyThreads.tmapreduce( + vcat, + OhMyThreads.chunks(eachindex(_init, run_seeds); n=nchunks); + chunking=false, + ) do chunk + chunk_rng = take!(rng_pool) + # `Optim` optimizers may carry mutable state, so each task gets its own copy. + chunk_optimizer = deepcopy(optimizer) + try + return map(chunk) do i + Random.seed!(chunk_rng, run_seeds[i]) + result = pathfinder( + optim_fun; + rng=chunk_rng, + history_length, + optimizer=chunk_optimizer, + ndraws=ndraws_per_run, + init=_init[i], + ntasks=ntasks_per_run, + ndraws_elbo, + kwargs..., + ) + put!(progress_channel, true) + yield() + return result + end + finally + put!(rng_pool, chunk_rng) + end + end + else + map(_init, run_seeds) do init, seed + Random.seed!(copy_rng, seed) + result = pathfinder( + optim_fun; + rng=copy_rng, + history_length, + optimizer=optimizer, + ndraws=ndraws_per_run, + init, + ntasks=ntasks_per_run, + ndraws_elbo, + kwargs..., + ) + put!(progress_channel, true) + yield() + return result + end + end + finally + put!(progress_channel, false) + close(progress_channel) + wait(progress_taskref[]) + end end - iter_optimizers = fill(optimizer, nruns) - iter_sp = - if executor isa Folds.ThreadedEx - # temporary workaround due to - # https://github.com/JuliaFolds2/Transducers.jl/issues/10 - # also support optimizers that store state - zip(_init, Iterators.map(deepcopy, iter_optimizers)) - else - Transducers.withprogress(zip(_init, iter_optimizers); interval=1e-3) - end |> trans - pathfinder_results = Folds.collect(iter_sp, executor) - fit_distributions = - pathfinder_results |> Transducers.Map(x -> x.fit_distribution) |> collect - draws_all = reduce(hcat, pathfinder_results |> Transducers.Map(x -> x.draws)) + fit_distributions = map(x -> x.fit_distribution, pathfinder_results) + draws_all = mapreduce(x -> x.draws, hcat, pathfinder_results) # draw samples from augmented mixture model inds = axes(draws_all, 2) sample_inds, psis_result = if importance - log_densities_fit = - pathfinder_results |> - Transducers.MapCat() do x - return Distributions.logpdf(x.fit_distribution, x.draws) - end |> - collect - iter_logp = eachcol(draws_all) |> Transducers.Map(logp) - log_densities_target = Folds.collect(iter_logp, executor) + log_densities_fit = if threaded + OhMyThreads.tmapreduce( + x -> Distributions.logpdf(x.fit_distribution, x.draws), + vcat, + pathfinder_results; + nchunks=nchunks, + ) + else + mapreduce( + x -> Distributions.logpdf(x.fit_distribution, x.draws), + vcat, + pathfinder_results, + ) + end + log_densities_target = if threaded + OhMyThreads.tmap(logp, eachcol(draws_all); nchunks=nchunks) + else + map(logp, eachcol(draws_all)) + end log_densities_ratios = log_densities_target - log_densities_fit resample(rng, inds, log_densities_ratios, ndraws) else diff --git a/src/mvnormal.jl b/src/mvnormal.jl index c89f978c9..8914b2b83 100644 --- a/src/mvnormal.jl +++ b/src/mvnormal.jl @@ -13,12 +13,10 @@ were rejected due to keeping the inverse Hessian positive definite. """ function fit_mvnormals(θs, ∇logpθs; kwargs...) Σs, num_bfgs_updates_rejected = lbfgs_inverse_hessians(θs, ∇logpθs; kwargs...) - trans = Transducers.MapSplat() do Σ, ∇logpθ, θ + dists = map(Σs, ∇logpθs, θs) do Σ, ∇logpθ, θ μ = muladd(Σ, ∇logpθ, θ) return Distributions.MvNormal(μ, Σ) end - l = length(Σs) - dists = @views(zip(Σs, ∇logpθs[1:l], θs[1:l])) |> trans |> collect return dists, num_bfgs_updates_rejected end diff --git a/src/singlepath.jl b/src/singlepath.jl index a2fc8bfe3..73ee98f6b 100644 --- a/src/singlepath.jl +++ b/src/singlepath.jl @@ -111,14 +111,10 @@ $(_ARGUMENT_DOCSTRING) - `ndraws_elbo::Int=$DEFAULT_NDRAWS_ELBO`: Number of draws used to estimate the ELBO - `ndraws::Int=ndraws_elbo`: number of approximate draws to return - `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. The default - ([`Transducers.SequentialEx()`](@extref `Transducers.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()`](@extref `Transducers.PreferParallel`) may be used to - parallelize log-density evaluation. This is generally only faster for expensive log - density functions. +- `ntasks::Int=1`: maximum number of parallel tasks used to compute ELBO estimates across + the fitted distributions. The default `ntasks = 1` runs sequentially; larger values + parallelize the ELBO evaluation, in which case the log-density function must be + thread-safe. Results are reproducible regardless of `ntasks`. - `history_length::Int=$DEFAULT_HISTORY_LENGTH`: Size of the history used to approximate the inverse Hessian. - `optimizer`: Optimizer to be used for constructing trajectory. Can be any optimizer @@ -293,7 +289,7 @@ function _pathfinder( history_length::Int=DEFAULT_HISTORY_LENGTH, optimizer=default_optimizer(history_length), ndraws_elbo=DEFAULT_NDRAWS_ELBO, - executor::Transducers.Executor=Transducers.SequentialEx(), + ntasks::Int=1, kwargs..., ) # compute trajectory @@ -308,7 +304,7 @@ function _pathfinder( # find ELBO-maximizing distribution fit_iteration, elbo_estimates = @views maximize_elbo( - rng, logp, fit_distributions[(begin + 1):end], ndraws_elbo, executor + rng, logp, fit_distributions[(begin + 1):end], ndraws_elbo, ntasks ) if isempty(elbo_estimates) success = false diff --git a/src/transducers.jl b/src/transducers.jl deleted file mode 100644 index 953e98845..000000000 --- a/src/transducers.jl +++ /dev/null @@ -1,18 +0,0 @@ -# transducer-friendly findmax, ignoring NaNs -function _findmax(x) - return Transducers.foldxl(x |> Transducers.Enumerate(); init=missing) do xmax_imax, i_xi - xmax_imax === missing && return reverse(i_xi) - i, xi = i_xi - isnan(xi) && return xmax_imax - xmax = first(xmax_imax) - isnan(xmax) && return (xi, i) - return xi > xmax ? (xi, i) : xmax_imax - end -end - -if !hasmethod(size, Tuple{Transducers.ProgressLoggingFoldable}) - # WARNING: Type piracy! - # https://github.com/JuliaFolds/Transducers.jl/issues/521 - # this method is necessary for Transducers version earlier than v0.4.82 - Base.size(x::Transducers.ProgressLoggingFoldable) = size(x.foldable) -end diff --git a/src/utils.jl b/src/utils.jl new file mode 100644 index 000000000..36e3a465b --- /dev/null +++ b/src/utils.jl @@ -0,0 +1,24 @@ +# Internal helpers + +# Return `(maxvalue, argmaxindex)` of `key.(xs)`, skipping `NaN`s. If the +# first observed value is `NaN`, it is retained until a non-`NaN` entry +# replaces it; if every entry is `NaN`, the first one is returned (matching +# the semantics the success check in `singlepath.jl` relies on). +_findmax_skipnan(xs) = _findmax_skipnan(identity, xs) + +function _findmax_skipnan(key, xs) + state = missing + for (i, x) in enumerate(xs) + xi = key(x) + if state === missing + state = (xi, i) + continue + end + isnan(xi) && continue + xmax = first(state) + if isnan(xmax) || xi > xmax + state = (xi, i) + end + end + return state +end diff --git a/test/Project.toml b/test/Project.toml index 50697402b..7ce2214bb 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -16,7 +16,6 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" -Transducers = "28d57a85-8fef-5791-bfe6-a80928e7c999" [compat] ADTypes = "0.2.5, 1" @@ -30,13 +29,12 @@ OptimizationBase = "0.0.5, 1, 2, 3, 4, 5.0" OptimizationNLopt = "0.2, 0.3" PDMats = "0.11.35" PSIS = "0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9" -Pathfinder = "0.9.27" +Pathfinder = "0.10" ProgressLogging = "0.1.4" Random = "1" ReverseDiff = "1.15" SciMLBase = "2.30" Test = "1" -Transducers = "0.4.81" [sources] Pathfinder = {path = ".."} diff --git a/test/elbo.jl b/test/elbo.jl index 4faf2f7f3..631af75ad 100644 --- a/test/elbo.jl +++ b/test/elbo.jl @@ -2,7 +2,6 @@ using Distributions using Pathfinder using Random using Test -using Transducers @testset "ELBO estimation" begin @testset "elbo_and_samples" begin @@ -34,22 +33,17 @@ 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) - executors = [SequentialEx(), ThreadedEx()] - @testset "$executor" for executor in executors + @testset "ntasks=$ntasks" for ntasks in (1, Threads.nthreads()) rng = Random.seed!(Random.default_rng(), 42) - lopt, estimates = @inferred Pathfinder.maximize_elbo( - rng, logp, dists, 100, executor - ) + lopt, estimates = Pathfinder.maximize_elbo(rng, logp, dists, 100, ntasks) @test lopt == 3 @test estimates[lopt].value ≈ 0 rng = Random.seed!(Random.default_rng(), 42) - lopt2, estimates2 = Pathfinder.maximize_elbo(rng, logp, dists, 100, executor) + lopt2, estimates2 = Pathfinder.maximize_elbo(rng, logp, dists, 100, ntasks) @test lopt2 == lopt @test getproperty.(estimates2, :value) == getproperty.(estimates, :value) @test getproperty.(estimates2, :std_err) == getproperty.(estimates, :std_err) - lopt3, estimates3 = @inferred Pathfinder.maximize_elbo( - rng, logp, dists[2:1], 100, executor - ) + lopt3, estimates3 = Pathfinder.maximize_elbo(rng, logp, dists[2:1], 100, ntasks) @test lopt3 == 0 @test isempty(estimates3) end diff --git a/test/multipath.jl b/test/multipath.jl index 7088be772..2a4e7493b 100644 --- a/test/multipath.jl +++ b/test/multipath.jl @@ -8,7 +8,6 @@ using PSIS using ReverseDiff using SciMLBase using Test -using Transducers @testset "multi path pathfinder" begin @testset "MvNormal" begin @@ -24,12 +23,10 @@ using Transducers ℓ = build_logdensityproblem(logp, dim, 2) rngs = [MersenneTwister(), Random.default_rng()] seed = 76 - @testset for rng in rngs - executor = rng isa MersenneTwister ? SequentialEx() : ThreadedEx() - + @testset for rng in rngs, ntasks in (1, Threads.nthreads()) Random.seed!(rng, seed) result = multipathfinder( - ℓ, ndraws; nruns, ndraws_elbo, ndraws_per_run, rng, executor + ℓ, ndraws; nruns, ndraws_elbo, ndraws_per_run, rng, ntasks ) @test result isa MultiPathfinderResult @test result.input === ℓ @@ -66,7 +63,7 @@ using Transducers Random.seed!(rng, seed) result2 = multipathfinder( - ℓ, ndraws; nruns, ndraws_elbo, ndraws_per_run, rng, executor + ℓ, ndraws; nruns, ndraws_elbo, ndraws_per_run, rng, ntasks ) @test result2.fit_distribution == result.fit_distribution @test result2.draws == result.draws @@ -74,7 +71,7 @@ using Transducers Random.seed!(rng, seed) result3 = multipathfinder( - ℓ, ndraws; nruns, ndraws_elbo, ndraws_per_run, rng, executor + ℓ, ndraws; nruns, ndraws_elbo, ndraws_per_run, rng, ntasks ) for (c1, c2) in zip(result.fit_distribution.components, result3.fit_distribution.components) @@ -107,4 +104,43 @@ using Transducers @test_throws ArgumentError multipathfinder(ℓ, 10; nruns=0) multipathfinder(ℓ, 10; nruns=2) end + + @testset "reproducibility across ntasks" begin + logp(x) = -sum(abs2, x) / 2 + ℓ = build_logdensityproblem(logp, 5, 2) + seed = 19 + nruns = 6 + ndraws = 200 + nthreads = Threads.nthreads() + + @testset "explicit rng" begin + rng = MersenneTwister(seed) + serial = multipathfinder(ℓ, ndraws; rng, nruns, ntasks=1, ntasks_per_run=1) + Random.seed!(rng, seed) + threaded = multipathfinder( + ℓ, ndraws; rng, nruns, ntasks=nthreads, ntasks_per_run=nthreads + ) + @test serial.draws == threaded.draws + @test serial.draw_component_ids == threaded.draw_component_ids + @test [c.μ for c in serial.fit_distribution.components] == + [c.μ for c in threaded.fit_distribution.components] + @test [c.Σ for c in serial.fit_distribution.components] == + [c.Σ for c in threaded.fit_distribution.components] + end + + @testset "default rng" begin + Random.seed!(seed) + serial = multipathfinder(ℓ, ndraws; nruns, ntasks=1, ntasks_per_run=1) + Random.seed!(seed) + threaded = multipathfinder( + ℓ, ndraws; nruns, ntasks=nthreads, ntasks_per_run=nthreads + ) + @test serial.draws == threaded.draws + @test serial.draw_component_ids == threaded.draw_component_ids + @test [c.μ for c in serial.fit_distribution.components] == + [c.μ for c in threaded.fit_distribution.components] + @test [c.Σ for c in serial.fit_distribution.components] == + [c.Σ for c in threaded.fit_distribution.components] + end + end end diff --git a/test/runtests.jl b/test/runtests.jl index b03bccbf2..f01472419 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -6,7 +6,7 @@ Random.seed!(0) @testset "Pathfinder.jl" begin include("test_utils.jl") - include("transducers.jl") + include("utils.jl") include("woodbury.jl") include("optimize.jl") include("inverse_hessian.jl") diff --git a/test/singlepath.jl b/test/singlepath.jl index e8f7d69c7..d77c8fab5 100644 --- a/test/singlepath.jl +++ b/test/singlepath.jl @@ -9,7 +9,6 @@ using Random using ReverseDiff using SciMLBase using Test -using Transducers @testset "single path pathfinder" begin @testset "IsoNormal" begin @@ -18,14 +17,16 @@ using Transducers ndraws = 100 rngs = [MersenneTwister(), Random.default_rng()] seed = 42 - @testset for dim in [1, 5, 10, 100], rng in rngs - executor = rng isa MersenneTwister ? SequentialEx() : ThreadedEx() + @testset for dim in [1, 5, 10, 100], + rng in rngs, + ntasks in (1, Threads.nthreads()) + ℓ = build_logdensityproblem(logp, 5, 2) init = randn(dim) Random.seed!(rng, seed) # less restrictive type check to work around https://github.com/mlcolab/Pathfinder.jl/issues/142 # TODO: remove this workaround once the issue is fixed - result = @inferred PathfinderResult pathfinder(ℓ; init, ndraws, rng, executor) + result = @inferred PathfinderResult pathfinder(ℓ; init, ndraws, rng, ntasks) @test result isa PathfinderResult @test result.input === ℓ @test result.optim_prob isa SciMLBase.OptimizationProblem @@ -53,14 +54,14 @@ using Transducers argmax(getproperty.(result.elbo_estimates, :value)) Random.seed!(rng, seed) - result2 = pathfinder(ℓ; init, ndraws, rng, executor) + result2 = pathfinder(ℓ; init, ndraws, rng, ntasks) @test result2.fit_iteration == result.fit_iteration @test result2.draws == result.draws @test getproperty.(result2.elbo_estimates, :value) == getproperty.(result.elbo_estimates, :value) ndraws = 2 - result3 = pathfinder(ℓ; init, ndraws, executor) + result3 = pathfinder(ℓ; init, ndraws, ntasks) @test size(result3.draws) == (dim, ndraws) end end @@ -78,25 +79,23 @@ using Transducers logp(x) = -dot(x, P, x) / 2 dim = 5 ℓ = build_logdensityproblem(logp, dim, 2) - ndraws_elbo = 100 + ndraws_elbo = 500 rngs = [MersenneTwister(), Random.default_rng()] x = randn(dim) seed = 38 optimizer = Optim.LBFGS(; m=6) - @testset for rng in rngs - executor = rng isa MersenneTwister ? SequentialEx() : ThreadedEx() - + @testset for rng in rngs, ntasks in (1, Threads.nthreads()) Random.seed!(rng, seed) # less restrictive type check to work around https://github.com/mlcolab/Pathfinder.jl/issues/142 # TODO: remove this workaround once the issue is fixed result = @inferred PathfinderResult pathfinder( - ℓ; rng, optimizer, ndraws_elbo, executor + ℓ; rng, optimizer, ndraws_elbo, ntasks ) @test result.input === ℓ @test result.fit_distribution.Σ ≈ Σ rtol = 1e-1 @test result.optimizer == optimizer Random.seed!(rng, seed) - result2 = pathfinder(ℓ; rng, optimizer, ndraws_elbo, executor) + result2 = pathfinder(ℓ; rng, optimizer, ndraws_elbo, ntasks) @test result2.fit_distribution == result.fit_distribution @test result2.draws == result.draws @test getproperty.(result2.elbo_estimates, :value) == @@ -171,4 +170,36 @@ using Transducers @test_throws ArgumentError pathfinder(build_logdensityproblem(logp, 0, 2)) pathfinder(build_logdensityproblem(logp, 3, 2)) end + + @testset "reproducibility across ntasks" begin + logp(x) = -sum(abs2, x) / 2 + ℓ = build_logdensityproblem(logp, 5, 2) + seed = 17 + + @testset "explicit rng" begin + rng = MersenneTwister(seed) + serial = pathfinder(ℓ; rng, ndraws=100, ntasks=1) + Random.seed!(rng, seed) + threaded = pathfinder(ℓ; rng, ndraws=100, ntasks=Threads.nthreads()) + @test serial.draws == threaded.draws + @test serial.fit_iteration == threaded.fit_iteration + @test getproperty.(serial.elbo_estimates, :value) == + getproperty.(threaded.elbo_estimates, :value) + @test serial.fit_distribution.μ == threaded.fit_distribution.μ + @test serial.fit_distribution.Σ == threaded.fit_distribution.Σ + end + + @testset "default rng" begin + Random.seed!(seed) + serial = pathfinder(ℓ; ndraws=100, ntasks=1) + Random.seed!(seed) + threaded = pathfinder(ℓ; ndraws=100, ntasks=Threads.nthreads()) + @test serial.draws == threaded.draws + @test serial.fit_iteration == threaded.fit_iteration + @test getproperty.(serial.elbo_estimates, :value) == + getproperty.(threaded.elbo_estimates, :value) + @test serial.fit_distribution.μ == threaded.fit_distribution.μ + @test serial.fit_distribution.Σ == threaded.fit_distribution.Σ + end + end end diff --git a/test/transducers.jl b/test/transducers.jl deleted file mode 100644 index cd78c8f8c..000000000 --- a/test/transducers.jl +++ /dev/null @@ -1,15 +0,0 @@ -using Pathfinder -using Random -using Test -using Transducers - -@testset "transducers integration" begin - @testset "_findmax" begin - x = randn(100) - @test Pathfinder._findmax(x) == findmax(x) - @test Pathfinder._findmax(x |> Map(sin)) == findmax(sin.(x)) - @test Pathfinder._findmax([NaN, 3.0, 1.0]) === (3.0, 2) - @test Pathfinder._findmax([NaN, NaN, NaN]) === (NaN, 1) - @test Pathfinder._findmax([2.0, NaN, 4.0]) === (4.0, 3) - end -end diff --git a/test/utils.jl b/test/utils.jl new file mode 100644 index 000000000..418d487aa --- /dev/null +++ b/test/utils.jl @@ -0,0 +1,14 @@ +using Pathfinder +using Random +using Test + +@testset "parallel helpers" begin + @testset "_findmax_skipnan" begin + x = randn(100) + @test Pathfinder._findmax_skipnan(x) == findmax(x) + @test Pathfinder._findmax_skipnan(sin, x) == findmax(sin.(x)) + @test Pathfinder._findmax_skipnan([NaN, 3.0, 1.0]) === (3.0, 2) + @test Pathfinder._findmax_skipnan([NaN, NaN, NaN]) === (NaN, 1) + @test Pathfinder._findmax_skipnan([2.0, NaN, 4.0]) === (4.0, 3) + end +end From b1d5e8615443c391f00a41e91c99aeafd0fa6012 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Mon, 25 May 2026 11:47:07 +0200 Subject: [PATCH 2/7] refactor: drop Optimization umbrella in favor of OptimizationBase `Optimization.solve` is the same `CommonSolve.solve` exposed by `SciMLBase` (which Pathfinder already depends on), and `OptimizationState` is now defined directly in `OptimizationBase 5`. Switching to `SciMLBase.solve` and importing `OptimizationState` from `OptimizationBase` lets us drop the `Optimization` umbrella and its `AbstractTrees`/`ProgressMeter`/`TerminalLoggers`/etc. tail. Compat floors for `Optim`, `OptimizationBase`, and `OptimizationOptimJL` are tightened so the declared ranges actually intersect. Co-Authored-By: Claude Opus 4.7 --- Project.toml | 8 +++----- src/Pathfinder.jl | 8 +------- src/optimize.jl | 2 +- test/Project.toml | 6 ++---- test/multipath.jl | 1 - test/optimize.jl | 8 ++------ test/singlepath.jl | 1 - 7 files changed, 9 insertions(+), 25 deletions(-) diff --git a/Project.toml b/Project.toml index 5ac6670cb..bfe7e7981 100644 --- a/Project.toml +++ b/Project.toml @@ -15,7 +15,6 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c" OhMyThreads = "67456a42-1dca-4109-a031-0a68de7e3ad5" Optim = "429524aa-4258-5aef-a3af-852621145aeb" -Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba" OptimizationBase = "bca83a33-5cc9-4baa-983d-23429ab6bcbb" OptimizationOptimJL = "36348300-93cb-4f02-beb5-3c3902f8871e" PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150" @@ -52,10 +51,9 @@ LinearAlgebra = "1" LogDensityProblems = "2.1.0" MCMCChains = "6.0.2, 7" OhMyThreads = "0.8.5" -Optim = "1.7.2, 2.0" -Optimization = "3.22, 4, 5" -OptimizationBase = "0.0.5, 1, 2, 3, 4, 5.0" -OptimizationOptimJL = "0.2.1, 0.3, 0.4" +Optim = "2" +OptimizationBase = "5" +OptimizationOptimJL = "0.4.10" PDMats = "0.11.35" PSIS = "0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9" ProgressLogging = "0.1.4" diff --git a/src/Pathfinder.jl b/src/Pathfinder.jl index f71d1be67..1f66811d2 100644 --- a/src/Pathfinder.jl +++ b/src/Pathfinder.jl @@ -7,7 +7,7 @@ using LinearAlgebra using LogDensityProblems: LogDensityProblems using OhMyThreads: OhMyThreads using Optim: Optim, LineSearches -using Optimization: Optimization +using OptimizationBase: OptimizationBase, OptimizationState using OptimizationOptimJL: OptimizationOptimJL using PDMats: PDMats using ProgressLogging: ProgressLogging @@ -18,12 +18,6 @@ using SciMLBase: SciMLBase using Statistics: Statistics using StatsBase: StatsBase -if isdefined(Optimization, :OptimizationState) - using Optimization: OptimizationState -else - using OptimizationBase: OptimizationState -end - # Declare and export the public API export PathfinderResult, MultiPathfinderResult export pathfinder, multipathfinder diff --git a/src/optimize.jl b/src/optimize.jl index f1d658822..9d7515c8f 100644 --- a/src/optimize.jl +++ b/src/optimize.jl @@ -51,7 +51,7 @@ function optimize_with_trace( _callback = OptimizationCallback( xs, fxs, ∇fxs, progress_name, progress_id, maxiters, callback, fail_on_nonfinite ) - sol = Optimization.solve(prob, optimizer; callback=_callback, maxiters, kwargs...) + sol = SciMLBase.solve(prob, optimizer; callback=_callback, maxiters, kwargs...) _∇fxs = _fill_missing_gradient_values!(∇fxs, xs, sol.cache.f) diff --git a/test/Project.toml b/test/Project.toml index 7ce2214bb..b9782685e 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -5,7 +5,6 @@ ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c" Optim = "429524aa-4258-5aef-a3af-852621145aeb" -Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba" OptimizationBase = "bca83a33-5cc9-4baa-983d-23429ab6bcbb" OptimizationNLopt = "4e6fcdb7-1186-4e1f-a706-475e75c168bb" PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150" @@ -24,9 +23,8 @@ ForwardDiff = "0.10.26, 1" LinearAlgebra = "1" LogDensityProblems = "2.1.0" Optim = "1.7.2, 2.0" -Optimization = "3.22, 4, 5" -OptimizationBase = "0.0.5, 1, 2, 3, 4, 5.0" -OptimizationNLopt = "0.2, 0.3" +OptimizationBase = "5" +OptimizationNLopt = "0.3" PDMats = "0.11.35" PSIS = "0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9" Pathfinder = "0.10" diff --git a/test/multipath.jl b/test/multipath.jl index 2a4e7493b..97e899b67 100644 --- a/test/multipath.jl +++ b/test/multipath.jl @@ -2,7 +2,6 @@ using ADTypes using Distributions using ForwardDiff using LinearAlgebra -using Optimization using Pathfinder using PSIS using ReverseDiff diff --git a/test/optimize.jl b/test/optimize.jl index 6f7a6c5bf..493d5e26c 100644 --- a/test/optimize.jl +++ b/test/optimize.jl @@ -1,18 +1,14 @@ using ADTypes using LinearAlgebra using Optim -using Optimization using OptimizationNLopt using Pathfinder using ProgressLogging using ReverseDiff using SciMLBase using Test -if isdefined(Optimization, :OptimizationState) - using Optimization: OptimizationState -else - using OptimizationBase: OptimizationState -end + +using OptimizationBase: OptimizationState @testset "build_optim_function" begin n = 20 diff --git a/test/singlepath.jl b/test/singlepath.jl index d77c8fab5..d9e99af6e 100644 --- a/test/singlepath.jl +++ b/test/singlepath.jl @@ -3,7 +3,6 @@ using Distributions using ForwardDiff using LinearAlgebra using Optim -using Optimization using Pathfinder using Random using ReverseDiff From 83df232c7f91296f80f6f4584793d3e566bfd7fc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Mon, 25 May 2026 12:11:43 +0200 Subject: [PATCH 3/7] Update environments of integration tests --- test/integration/AdvancedHMC/Project.toml | 4 ++-- test/integration/DynamicHMC/Project.toml | 4 ++-- test/integration/Turing/Project.toml | 2 +- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/test/integration/AdvancedHMC/Project.toml b/test/integration/AdvancedHMC/Project.toml index b0b1dca74..0322c5b7d 100644 --- a/test/integration/AdvancedHMC/Project.toml +++ b/test/integration/AdvancedHMC/Project.toml @@ -26,8 +26,8 @@ LinearAlgebra = "1" LogDensityProblems = "2.1.0" LogDensityProblemsAD = "1.7" MCMCDiagnosticTools = "0.3" -Optim = "1.7.2, 2.0" -Pathfinder = "0.9.21" +Optim = "2" +Pathfinder = "0.10" Random = "1" Statistics = "1" StatsFuns = "1" diff --git a/test/integration/DynamicHMC/Project.toml b/test/integration/DynamicHMC/Project.toml index b29bca51f..cae2bea78 100644 --- a/test/integration/DynamicHMC/Project.toml +++ b/test/integration/DynamicHMC/Project.toml @@ -22,8 +22,8 @@ LinearAlgebra = "1" LogDensityProblems = "2.1.0" LogDensityProblemsAD = "1.7" MCMCDiagnosticTools = "0.3" -Optim = "1.7.2, 2.0" -Pathfinder = "0.9.18" +Optim = "2" +Pathfinder = "0.10" Random = "1" Statistics = "1" StatsFuns = "1" diff --git a/test/integration/Turing/Project.toml b/test/integration/Turing/Project.toml index 6d6713a85..ad32d939a 100644 --- a/test/integration/Turing/Project.toml +++ b/test/integration/Turing/Project.toml @@ -21,7 +21,7 @@ DynamicPPL = "0.40.6" FlexiChains = "0.3.4" LinearAlgebra = "1" LogDensityProblems = "2.1.0" -Pathfinder = "0.9.27" +Pathfinder = "0.10" Random = "1" ReverseDiff = "1.15" Test = "1" From 749ef2207940a4f8be796103f23bacf775031c1b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Mon, 25 May 2026 12:12:35 +0200 Subject: [PATCH 4/7] Update docs environment --- docs/Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/Project.toml b/docs/Project.toml index f174f0d0a..dd56d3e3e 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -39,7 +39,7 @@ ForwardDiff = "0.10.26, 1" LinearAlgebra = "1" LogDensityProblems = "2.1.0" LogDensityProblemsAD = "1.7" -Pathfinder = "0.9.27" +Pathfinder = "0.10" Printf = "1" ReverseDiff = "1.15" StableRNGs = "1" From 4270ee783a73b6c4f3cdc7bce3c1e5983240ed2b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Mon, 25 May 2026 14:45:42 +0200 Subject: [PATCH 5/7] Change seed of Turing example --- ext/PathfinderTuringExt.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ext/PathfinderTuringExt.jl b/ext/PathfinderTuringExt.jl index cf4db8931..3ed77fa3e 100644 --- a/ext/PathfinderTuringExt.jl +++ b/ext/PathfinderTuringExt.jl @@ -214,7 +214,7 @@ Run multi-path Pathfinder on a Turing `model`. ```jldoctest julia> using Pathfinder, Turing, StableRNGs -julia> rng = StableRNG(42); +julia> rng = StableRNG(3); julia> @model function demo_model() α ~ Normal(0, 1) From c4da1f6a3b507d067ec56578b292543987c43b27 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Mon, 15 Jun 2026 16:59:16 +0200 Subject: [PATCH 6/7] refactor: extract shared parallel-map helpers to reduce duplication Both maximize_elbo and multipathfinder open-coded the same reproducible, seeded, chunked parallel-map idiom (RNG-pool Channel + tmapreduce over chunks + sequential fallback), and the importance resampling repeated a maybe-parallel map/reduce twice. Add internal helpers to src/utils.jl: - _nchunks: single source of truth for chunk/task count from ntasks - _chunk_tmap: generic chunked parallel map over equal-axes arrays with per-chunk setup state; seeding now lives at the call sites - _maybe_tmap / _maybe_tmapreduce: rng-free resampling reductions The RNG pool is replaced by a per-chunk copy(rng) built in setup; this is result-identical since every chunk RNG is re-seeded per element before use. Per-chunk optimizer deepcopy is preserved via setup. Algorithm logic and reproducibility across ntasks are unchanged. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/elbo.jl | 32 ++------------- src/multipath.jl | 104 ++++++++++++----------------------------------- src/utils.jl | 48 ++++++++++++++++++++++ test/utils.jl | 59 +++++++++++++++++++++++++++ 4 files changed, 137 insertions(+), 106 deletions(-) diff --git a/src/elbo.jl b/src/elbo.jl index 727496308..e5f08726a 100644 --- a/src/elbo.jl +++ b/src/elbo.jl @@ -1,34 +1,8 @@ function maximize_elbo(rng, logp, dists, ndraws, ntasks::Int) seeds = rand!(rng, similar(dists, UInt64)) - copy_rng = copy(rng) - nchunks = min(length(dists), ntasks) - estimates = if nchunks > 1 && Threads.nthreads() > 1 - nrngs = min(nchunks, Threads.nthreads()) - rng_pool = Channel{typeof(copy_rng)}(nrngs) - put!(rng_pool, copy_rng) - for _ in 2:nrngs - put!(rng_pool, copy(rng)) - end - OhMyThreads.tmapreduce( - vcat, - OhMyThreads.chunks(eachindex(dists, seeds); n=nchunks); - chunking=false, - ) do chunk - chunk_rng = take!(rng_pool) - try - return map(chunk) do i - Random.seed!(chunk_rng, seeds[i]) - return elbo_and_samples(chunk_rng, logp, dists[i], ndraws) - end - finally - put!(rng_pool, chunk_rng) - end - end - else - map(dists, seeds) do dist, seed - Random.seed!(copy_rng, seed) - return elbo_and_samples(copy_rng, logp, dist, ndraws) - end + estimates = _chunk_tmap(dists, seeds; ntasks, setup=() -> copy(rng)) do chunk_rng, dist, seed + Random.seed!(chunk_rng, seed) + return elbo_and_samples(chunk_rng, logp, dist, ndraws) end isempty(estimates) && return 0, estimates _, iteration_opt = _findmax_skipnan(est -> est.value, estimates) diff --git a/src/multipath.jl b/src/multipath.jl index 0d3df3107..f7d9a12e0 100644 --- a/src/multipath.jl +++ b/src/multipath.jl @@ -160,15 +160,12 @@ function multipathfinder( # run pathfinder independently from each starting point run_seeds = rand!(rng, similar(_init, UInt64)) - copy_rng = copy(rng) - nchunks = min(nruns, ntasks) - threaded = nchunks > 1 && Threads.nthreads() > 1 pathfinder_results = ProgressLogging.@withprogress name = "Multi-path Pathfinder" begin progress_taskref = Ref{Task}() progress_channel = Channel{Bool}( min(nruns, 1_000); spawn=true, taskref=progress_taskref - ) do ch + ) do ch # Throttle progress logs as generally the logging system is not super performant: # - At most 1 progress log every 0.1 seconds # - At most 1 progress log every 0.5% progress @@ -189,60 +186,25 @@ function multipathfinder( end try - if threaded - rng_pool = Channel{typeof(copy_rng)}(nchunks) - put!(rng_pool, copy_rng) - for _ in 2:nchunks - put!(rng_pool, copy(rng)) - end - OhMyThreads.tmapreduce( - vcat, - OhMyThreads.chunks(eachindex(_init, run_seeds); n=nchunks); - chunking=false, - ) do chunk - chunk_rng = take!(rng_pool) - # `Optim` optimizers may carry mutable state, so each task gets its own copy. - chunk_optimizer = deepcopy(optimizer) - try - return map(chunk) do i - Random.seed!(chunk_rng, run_seeds[i]) - result = pathfinder( - optim_fun; - rng=chunk_rng, - history_length, - optimizer=chunk_optimizer, - ndraws=ndraws_per_run, - init=_init[i], - ntasks=ntasks_per_run, - ndraws_elbo, - kwargs..., - ) - put!(progress_channel, true) - yield() - return result - end - finally - put!(rng_pool, chunk_rng) - end - end - else - map(_init, run_seeds) do init, seed - Random.seed!(copy_rng, seed) - result = pathfinder( - optim_fun; - rng=copy_rng, - history_length, - optimizer=optimizer, - ndraws=ndraws_per_run, - init, - ntasks=ntasks_per_run, - ndraws_elbo, - kwargs..., - ) - put!(progress_channel, true) - yield() - return result - end + # `Optim` optimizers may carry mutable state, so each chunk gets its own copy. + _chunk_tmap( + _init, run_seeds; ntasks, setup=() -> (copy(rng), deepcopy(optimizer)) + ) do (chunk_rng, chunk_optimizer), init, seed + Random.seed!(chunk_rng, seed) + result = pathfinder( + optim_fun; + rng=chunk_rng, + history_length, + optimizer=chunk_optimizer, + ndraws=ndraws_per_run, + init, + ntasks=ntasks_per_run, + ndraws_elbo, + kwargs..., + ) + put!(progress_channel, true) + yield() + return result end finally put!(progress_channel, false) @@ -256,25 +218,13 @@ function multipathfinder( # draw samples from augmented mixture model inds = axes(draws_all, 2) sample_inds, psis_result = if importance - log_densities_fit = if threaded - OhMyThreads.tmapreduce( - x -> Distributions.logpdf(x.fit_distribution, x.draws), - vcat, - pathfinder_results; - nchunks=nchunks, - ) - else - mapreduce( - x -> Distributions.logpdf(x.fit_distribution, x.draws), - vcat, - pathfinder_results, - ) - end - log_densities_target = if threaded - OhMyThreads.tmap(logp, eachcol(draws_all); nchunks=nchunks) - else - map(logp, eachcol(draws_all)) - end + log_densities_fit = _maybe_tmapreduce( + x -> Distributions.logpdf(x.fit_distribution, x.draws), + vcat, + pathfinder_results, + ntasks, + ) + log_densities_target = _maybe_tmap(logp, eachcol(draws_all), ntasks) log_densities_ratios = log_densities_target - log_densities_fit resample(rng, inds, log_densities_ratios, ndraws) else diff --git a/src/utils.jl b/src/utils.jl index 36e3a465b..b6aa0ecde 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -1,5 +1,53 @@ # Internal helpers +# Number of parallel chunks for `len` items given `ntasks`; `1` runs sequentially. +function _nchunks(len::Int, ntasks::Int) + nchunks = min(len, ntasks) + if nchunks > 1 && Threads.nthreads() > 1 + return nchunks + else + return 1 + end +end + +# Map `f(s, xs[i]...)` over the elements of the equal-axes arrays `xss`, using at most +# `ntasks` parallel tasks and per-chunk state `s = setup()` built once per chunk. +function _chunk_tmap(f, xss::AbstractArray...; ntasks::Int, setup) + nchunks = _nchunks(length(first(xss)), ntasks) + if nchunks == 1 + let s = setup() + return map((xs...) -> f(s, xs...), xss...) + end + else + return OhMyThreads.tmapreduce( + vcat, OhMyThreads.chunks(eachindex(xss...); n=nchunks); chunking=false + ) do chunk + let s = setup() + return map(i -> f(s, map(xs -> xs[i], xss)...), chunk) + end + end + end +end + +# `map`/`mapreduce` using at most `ntasks` parallel tasks. +function _maybe_tmap(f, xs::AbstractArray, ntasks::Int) + nchunks = _nchunks(length(xs), ntasks) + if nchunks == 1 + return map(f, xs) + else + return OhMyThreads.tmap(f, xs; nchunks) + end +end + +function _maybe_tmapreduce(f, op, xs::AbstractArray, ntasks::Int) + nchunks = _nchunks(length(xs), ntasks) + if nchunks == 1 + return mapreduce(f, op, xs) + else + return OhMyThreads.tmapreduce(f, op, xs; nchunks) + end +end + # Return `(maxvalue, argmaxindex)` of `key.(xs)`, skipping `NaN`s. If the # first observed value is `NaN`, it is retained until a non-`NaN` entry # replaces it; if every entry is `NaN`, the first one is returned (matching diff --git a/test/utils.jl b/test/utils.jl index 418d487aa..9c0d93040 100644 --- a/test/utils.jl +++ b/test/utils.jl @@ -11,4 +11,63 @@ using Test @test Pathfinder._findmax_skipnan([NaN, NaN, NaN]) === (NaN, 1) @test Pathfinder._findmax_skipnan([2.0, NaN, 4.0]) === (4.0, 3) end + + @testset "_nchunks" begin + @test Pathfinder._nchunks(0, 4) == 1 + @test Pathfinder._nchunks(1, 4) == 1 + @test Pathfinder._nchunks(10, 1) == 1 + if Threads.nthreads() > 1 + @test Pathfinder._nchunks(10, 4) == 4 + @test Pathfinder._nchunks(2, 4) == 2 + end + end + + @testset "_maybe_tmap" begin + xs = randn(50) + @testset "ntasks=$ntasks" for ntasks in (1, Threads.nthreads()) + @test Pathfinder._maybe_tmap(sin, xs, ntasks) == map(sin, xs) + @test Pathfinder._maybe_tmap(sin, Float64[], ntasks) == Float64[] + end + end + + @testset "_maybe_tmapreduce" begin + xs = [randn(3) for _ in 1:20] + @testset "ntasks=$ntasks" for ntasks in (1, Threads.nthreads()) + @test Pathfinder._maybe_tmapreduce(identity, vcat, xs, ntasks) == + mapreduce(identity, vcat, xs) + # `+` on floats is non-associative, so parallel reduction order may differ + @test Pathfinder._maybe_tmapreduce(sum, +, xs, ntasks) ≈ mapreduce(sum, +, xs) + end + end + + @testset "_chunk_tmap" begin + xs = collect(1:50) + ys = randn(50) + @testset "ntasks=$ntasks" for ntasks in (1, Threads.nthreads()) + # matches sequential map, with correct multi-array pairing + @test Pathfinder._chunk_tmap(xs, ys; ntasks, setup=() -> 0) do _, x, y + return x + y + end == map(+, xs, ys) + # per-chunk state is passed through + @test Pathfinder._chunk_tmap(xs; ntasks, setup=() -> 10) do s, x + return s * x + end == map(x -> 10x, xs) + # empty arrays return empty + @test isempty(Pathfinder._chunk_tmap(Int[]; ntasks, setup=() -> 0) do _, x + return x + end) + end + + # the reproducible-seeding idiom yields identical output regardless of ntasks + function seeded(seed, n, ntasks) + rng = Random.seed!(Random.default_rng(), seed) + xs = collect(1:n) + seeds = rand!(rng, similar(xs, UInt64)) + return Pathfinder._chunk_tmap(xs, seeds; ntasks, setup=() -> copy(rng)) do r, _, s + Random.seed!(r, s) + return rand(r) + end + end + @test seeded(7, 50, 1) == seeded(7, 50, Threads.nthreads()) + end end From 6723ba3a2d2e84e861d51ac07c51b782027ead95 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Wed, 17 Jun 2026 12:17:30 +0200 Subject: [PATCH 7/7] Apply suggestions from code review Co-authored-by: Seth Axen --- test/elbo.jl | 8 ++++++-- test/multipath.jl | 2 +- test/singlepath.jl | 4 ++-- 3 files changed, 9 insertions(+), 5 deletions(-) diff --git a/test/elbo.jl b/test/elbo.jl index 631af75ad..d48bbf8cf 100644 --- a/test/elbo.jl +++ b/test/elbo.jl @@ -35,7 +35,9 @@ using Test dists = Normal.(0, σs) @testset "ntasks=$ntasks" for ntasks in (1, Threads.nthreads()) rng = Random.seed!(Random.default_rng(), 42) - lopt, estimates = Pathfinder.maximize_elbo(rng, logp, dists, 100, ntasks) + lopt, estimates = @inferred Pathfinder.maximize_elbo( + rng, logp, dists, 100, ntasks + ) @test lopt == 3 @test estimates[lopt].value ≈ 0 rng = Random.seed!(Random.default_rng(), 42) @@ -43,7 +45,9 @@ using Test @test lopt2 == lopt @test getproperty.(estimates2, :value) == getproperty.(estimates, :value) @test getproperty.(estimates2, :std_err) == getproperty.(estimates, :std_err) - lopt3, estimates3 = Pathfinder.maximize_elbo(rng, logp, dists[2:1], 100, ntasks) + lopt3, estimates3 = @inferred Pathfinder.maximize_elbo( + rng, logp, dists[2:1], 100, ntasks + ) @test lopt3 == 0 @test isempty(estimates3) end diff --git a/test/multipath.jl b/test/multipath.jl index 97e899b67..b4124045f 100644 --- a/test/multipath.jl +++ b/test/multipath.jl @@ -22,7 +22,7 @@ using Test ℓ = build_logdensityproblem(logp, dim, 2) rngs = [MersenneTwister(), Random.default_rng()] seed = 76 - @testset for rng in rngs, ntasks in (1, Threads.nthreads()) + @testset for rng in rngs, ntasks in unique((1, Threads.nthreads())) Random.seed!(rng, seed) result = multipathfinder( ℓ, ndraws; nruns, ndraws_elbo, ndraws_per_run, rng, ntasks diff --git a/test/singlepath.jl b/test/singlepath.jl index d9e99af6e..1dde6dd4e 100644 --- a/test/singlepath.jl +++ b/test/singlepath.jl @@ -18,7 +18,7 @@ using Test seed = 42 @testset for dim in [1, 5, 10, 100], rng in rngs, - ntasks in (1, Threads.nthreads()) + ntasks in unique((1, Threads.nthreads())) ℓ = build_logdensityproblem(logp, 5, 2) init = randn(dim) @@ -83,7 +83,7 @@ using Test x = randn(dim) seed = 38 optimizer = Optim.LBFGS(; m=6) - @testset for rng in rngs, ntasks in (1, Threads.nthreads()) + @testset for rng in rngs, ntasks in unique((1, Threads.nthreads())) Random.seed!(rng, seed) # less restrictive type check to work around https://github.com/mlcolab/Pathfinder.jl/issues/142 # TODO: remove this workaround once the issue is fixed