Merge remote-tracking branch 'origin/main' into breaking

Author: penelopeysm

.github/workflows/Tests.yml

@@ -26,16 +26,12 @@ jobs:
           # not included in the others.
           - name: "mcmc/gibbs"
             args: "mcmc/gibbs.jl"
-          - name: "mcmc/hmc"
-            args: "mcmc/hmc.jl"
-          - name: "mcmc/abstractmcmc"
-            args: "mcmc/abstractmcmc.jl"
           - name: "mcmc/Inference"
             args: "mcmc/Inference.jl"
-          - name: "mcmc/ess"
-            args: "mcmc/ess.jl"
+          - name: "ad"
+            args: "ad.jl"
           - name: "everything else"
-            args: "--skip mcmc/gibbs.jl mcmc/hmc.jl mcmc/abstractmcmc.jl mcmc/Inference.jl mcmc/ess.jl"
+            args: "--skip mcmc/gibbs.jl mcmc/Inference.jl ad.jl"
         runner:
           # Default
           - version: '1'

HISTORY.md

@@ -1,3 +1,8 @@
+# Release 0.38.4
+
+The minimum Julia version was increased to 1.10.2 (from 1.10.0).
+On versions before 1.10.2, `sample()` took an excessively long time to run (probably due to compilation).
+
 # Release 0.38.3
 
 `getparams(::Model, ::AbstractVarInfo)` now returns an empty `Float64[]` if the VarInfo contains no parameters.

Project.toml

@@ -83,7 +83,7 @@ Statistics = "1.6"
 StatsAPI = "1.6"
 StatsBase = "0.32, 0.33, 0.34"
 StatsFuns = "0.8, 0.9, 1"
-julia = "1.10"
+julia = "1.10.2"
 
 [extras]
 DynamicHMC = "bbc10e6e-7c05-544b-b16e-64fede858acb"

src/mcmc/hmc.jl

@@ -138,6 +138,36 @@ function AbstractMCMC.sample(
     end
 end
 
+function find_initial_params(
+    rng::Random.AbstractRNG,
+    model::DynamicPPL.Model,
+    varinfo::DynamicPPL.AbstractVarInfo,
+    hamiltonian::AHMC.Hamiltonian;
+    max_attempts::Int=1000,
+)
+    varinfo = deepcopy(varinfo)  # Don't mutate
+
+    for attempts in 1:max_attempts
+        theta = varinfo[:]
+        z = AHMC.phasepoint(rng, theta, hamiltonian)
+        isfinite(z) && return varinfo, z
+
+        attempts == 10 &&
+            @warn "failed to find valid initial parameters in $(attempts) tries; consider providing explicit initial parameters using the `initial_params` keyword"
+
+        # Resample and try again.
+        # NOTE: varinfo has to be linked to make sure this samples in unconstrained space
+        varinfo = last(
+            DynamicPPL.evaluate!!(model, rng, varinfo, DynamicPPL.SampleFromUniform())
+        )
+    end
+
+    # if we failed to find valid initial parameters, error
+    return error(
+        "failed to find valid initial parameters in $(max_attempts) tries. This may indicate an error with the model or AD backend; please open an issue at https://github.com/TuringLang/Turing.jl/issues",
+    )
+end
+
 function DynamicPPL.initialstep(
     rng::AbstractRNG,
     model::AbstractModel,
@@ -170,33 +200,14 @@ function DynamicPPL.initialstep(
     lp_grad_func = Base.Fix1(LogDensityProblems.logdensity_and_gradient, ldf)
     hamiltonian = AHMC.Hamiltonian(metric, lp_func, lp_grad_func)
 
-    # Compute phase point z.
-    z = AHMC.phasepoint(rng, theta, hamiltonian)
-
     # If no initial parameters are provided, resample until the log probability
-    # and its gradient are finite.
-    if initial_params === nothing
-        init_attempt_count = 1
-        while !isfinite(z)
-            if init_attempt_count == 10
-                @warn "failed to find valid initial parameters in $(init_attempt_count) tries; consider providing explicit initial parameters using the `initial_params` keyword"
-            end
-            if init_attempt_count == 1000
-                error(
-                    "failed to find valid initial parameters in $(init_attempt_count) tries. This may indicate an error with the model or AD backend; please open an issue at https://github.com/TuringLang/Turing.jl/issues",
-                )
-            end
-
-            # NOTE: This will sample in the unconstrained space.
-            vi = last(DynamicPPL.evaluate!!(model, rng, vi, SampleFromUniform()))
-            theta = vi[:]
-
-            hamiltonian = AHMC.Hamiltonian(metric, lp_func, lp_grad_func)
-            z = AHMC.phasepoint(rng, theta, hamiltonian)
-
-            init_attempt_count += 1
-        end
+    # and its gradient are finite. Otherwise, just use the existing parameters.
+    vi, z = if initial_params === nothing
+        find_initial_params(rng, model, vi, hamiltonian)
+    else
+        vi, AHMC.phasepoint(rng, theta, hamiltonian)
     end
+    theta = vi[:]
 
     # Cache current log density.
     log_density_old = getlogp(vi)

test/Project.toml

@@ -1,4 +1,5 @@
 [deps]
+ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
 AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf"
 AdvancedMH = "5b7e9947-ddc0-4b3f-9b55-0d8042f74170"
@@ -20,7 +21,6 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
 LogDensityProblemsAD = "996a588d-648d-4e1f-a8f0-a84b347e47b1"
 MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
-Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 NamedArrays = "86f7a689-2022-50b4-a561-43c23ac3c673"
 Optim = "429524aa-4258-5aef-a3af-852621145aeb"
 Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
@@ -39,6 +39,7 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
 
 [compat]
+ADTypes = "1"
 AbstractMCMC = "5"
 AbstractPPL = "0.9, 0.10, 0.11"
 AdvancedMH = "0.6, 0.7, 0.8"
@@ -52,15 +53,14 @@ Combinatorics = "1"
 Distributions = "0.25"
 DistributionsAD = "0.6.3"
 DynamicHMC = "2.1.6, 3.0"
-DynamicPPL = "0.36"
+DynamicPPL = "0.36.6"
 FiniteDifferences = "0.10.8, 0.11, 0.12"
 ForwardDiff = "0.10.12 - 0.10.32, 0.10"
 HypothesisTests = "0.11"
 LinearAlgebra = "1"
 LogDensityProblems = "2"
 LogDensityProblemsAD = "1.4"
 MCMCChains = "5, 6, 7"
-Mooncake = "0.4.95"
 NamedArrays = "0.9.4, 0.10"
 Optim = "1"
 Optimization = "3, 4"