Merge branch 'main' into qqy/sghmc

Author: yebai

.buildkite/pipeline.yml

@@ -0,0 +1,28 @@
+env:
+  # SECRET_CODECOV_TOKEN can be added here if needed for coverage reporting
+
+steps:
+  - label: "Julia v{{matrix.version}}, {{matrix.label}}"
+    plugins:
+      - JuliaCI/julia#v1:
+          version: "{{matrix.version}}"
+      # - JuliaCI/julia-coverage#v1:
+      #     dirs:
+      #       - src
+      #       - ext
+    command: julia --eval='println(pwd()); println(readdir()); include("test/CUDA/cuda.jl")'
+    agents:
+      queue: "juliagpu"
+      cuda: "*"
+    if: build.message !~ /\[skip tests\]/
+    timeout_in_minutes: 60
+    env:
+      LABEL: "{{matrix.label}}"
+      TEST_TYPE: ext
+    matrix:
+      setup:
+        version:
+          - "1"
+          - "1.10"
+        label:
+          - "cuda"

HISTORY.md

@@ -0,0 +1,46 @@
+# AdvancedHMC Changelog
+
+## v0.7.1
+
+  - README has been simplified, many docs transfered to docs: https://turinglang.org/AdvancedHMC.jl/dev/.
+  - ADTypes.jl can be used for specifying the AD backend in `Hamiltonian(metric, ℓπ, AutoForwardDiff())`.
+  - SimpleUnpack.jl and Requires.jl are removed from the dependency.
+  - `find_good_stepsize` now has fewer allocations.
+
+## v0.7.0
+
+  - Type piracies of Base.rand and Base.randn for vectors of RNGs are removed: Replace `rand(rngs::AbstractVector{<:Random.AbstractRNG})` with `map(rand, rngs)`, `randn(rngs::AbstractVector{<:Random.AbstractRNG})` with `map(randn, rngs)`, `rand(rngs::AbstractVector{<:Random.AbstractRNG}, T, n::Int) (for n == length(rngs))` with `map(Base.Fix2(rand, T), rngs)`, and `randn(rngs::AbstractVector{<:Random.AbstractRNG}, T, m::Int, n::Int) (for n == length(rngs))` with eg `reduce(hcat, map(rng -> randn(rng, T, m), rngs))`.
+  - Type piracy `Base.isfinite(x::AbstractVecOrMat)` is removed: Switch to `all(isfinite, x)` if you (possibly implicitly) relied on this definition
+  - Abstract fields of `NesterovDualAveraging`, `HMCDA`, `SliceTS`, and `MultinomialTS` are made concrete by adding type parameters: Update occurrences of these types (eg. in function signatures) if necessary
+  - Definitions of Base.rand for metrics are removed: Use the (internal) `AdvancedHMC.rand_momentum` function if you depend on this functionality and open an issue to further discuss the API
+
+## v0.5.0
+
+Convenience constructors for common samplers changed to:
+
+  - `HMC(leapfrog_stepsize::Real, n_leapfrog::Int)`
+  - `NUTS(target_acceptance::Real)`
+  - `HMCDA(target_acceptance::Real, integration_time::Real)`
+
+## v0.2.22
+
+Three functions are renamed.
+
+  - `Preconditioner(metric::AbstractMetric)` -> `MassMatrixAdaptor(metric)` and
+  - `NesterovDualAveraging(δ, integrator::AbstractIntegrator)` -> `StepSizeAdaptor(δ, integrator)`
+  - `find_good_eps` -> `find_good_stepsize`
+
+## v0.2.15
+
+`n_adapts` is no longer needed to construct `StanHMCAdaptor`; the old constructor is deprecated.
+
+## v0.2.8
+
+Two Hamiltonian trajectory sampling methods are renamed to avoid a name clash with Distributions.
+
+  - `Multinomial` -> `MultinomialTS`
+  - `Slice` -> `SliceTS`
+
+## v0.2.0
+
+The gradient function passed to `Hamiltonian` is supposed to return a value-gradient tuple now.

Project.toml

@@ -18,12 +18,14 @@ StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
 
 [weakdeps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
+ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 
 [extensions]
 AdvancedHMCADTypesExt = "ADTypes"
+AdvancedHMCComponentArraysExt = "ComponentArrays"
 AdvancedHMCCUDAExt = "CUDA"
 AdvancedHMCMCMCChainsExt = "MCMCChains"
 AdvancedHMCOrdinaryDiffEqExt = "OrdinaryDiffEq"
@@ -32,6 +34,7 @@ AdvancedHMCOrdinaryDiffEqExt = "OrdinaryDiffEq"
 ADTypes = "1"
 AbstractMCMC = "5.6"
 ArgCheck = "1, 2"
+ComponentArrays = "0.15"
 CUDA = "3, 4, 5"
 DocStringExtensions = "0.8, 0.9"
 LinearAlgebra = "<0.1, 1"

docs/make.jl

@@ -20,7 +20,6 @@ makedocs(;
         "Detailed API" => "api.md",
         "Interfaces" => "interfaces.md",
         "News" => "news.md",
-        "Change Log" => "changelog.md",
         "References" => "references.md",
     ],
 )

docs/src/autodiff.md

@@ -1,5 +1,11 @@
 # Gradient in AdvancedHMC.jl
 
-AdvancedHMC.jl supports automatic differentiation using [`LogDensityProblemsAD`](https://github.com/tpapp/LogDensityProblemsAD.jl) across various AD backends and allows user-specified gradients. While the default AD backend for AdvancedHMC.jl is ForwardDiff.jl, we can seamlessly change to other backend like Mooncake.jl using various syntax like `Hamiltonian(metric, ℓπ, AutoZygote())`. Different AD backend can also be pluged in using `Hamiltonian(metric, ℓπ, Zygote)`, `Hamiltonian(metric, ℓπ, Val(:Zygote))` but we recommend using ADTypes since that would allow you to have more freedom for specifying the AD backend.
+AdvancedHMC.jl supports automatic differentiation using [`LogDensityProblemsAD`](https://github.com/tpapp/LogDensityProblemsAD.jl) across various AD backends and allows user-specified gradients. While the default AD backend for AdvancedHMC.jl is ForwardDiff.jl, we can seamlessly change to other backend like Mooncake.jl using various syntax like `Hamiltonian(metric, ℓπ, AutoMooncake(; config = nothing))`. While some AD backends support syntax like `Hamiltonian(metric, ℓπ, Zygote)`, `Hamiltonian(metric, ℓπ, Val(:Zygote))`, we recommend using ADTypes since that would allow you to have more freedom for specifying the AD backend:
+
+```julia
+using AdvancedHMC, ADTypes, DifferentiationInterface, Mooncake, Zygote
+hamiltonian = Hamiltonian(metric, ℓπ, AutoMooncake(; config=nothing))
+hamiltonian = Hamiltonian(metric, ℓπ, AutoZygote())
+```
 
 In order to use user-specified gradients, please replace ForwardDiff.jl with `ℓπ_grad` in the `Hamiltonian` constructor as `Hamiltonian(metric, ℓπ, ℓπ_grad)`, where the gradient function `ℓπ_grad` should return a tuple containing both the log-posterior and its gradient, for example `ℓπ_grad(x) = (log_posterior, grad)`.

docs/src/changelog.md

@@ -0,0 +1,8 @@
+module AdvancedHMCComponentArraysExt
+
+using AdvancedHMC: AdvancedHMC, __axes
+using ComponentArrays: ComponentVecOrMat, getaxes
+
+AdvancedHMC.__axes(r::ComponentVecOrMat) = getaxes(r)
+
+end # module

ext/AdvancedHMCComponentArraysExt.jl

@@ -172,9 +172,12 @@ function Hamiltonian(metric::AbstractMetric, ℓπ, kind::Union{Symbol,Val,Modul
             ),
         )
     end
-    ℓ = LogDensityProblemsAD.ADgradient(
-        kind isa Val ? kind : Val(Symbol(kind)), ℓπ; kwargs...
-    )
+    _kind = if kind isa Val || kind isa Symbol
+        kind
+    else
+        Symbol(kind)
+    end
+    ℓ = LogDensityProblemsAD.ADgradient(_kind, ℓπ; kwargs...)
     return Hamiltonian(metric, ℓ)
 end
 

src/AdvancedHMC.jl

@@ -13,7 +13,7 @@ function adapt!(
     α::AbstractScalarOrVec{<:AbstractFloat},
     is_update::Bool=true,
 )
-    resize!(adaptor, θ)
+    resize_adaptor!(adaptor, size(θ))
     push!(adaptor, θ)
     is_update && update!(adaptor)
     return nothing
@@ -29,7 +29,7 @@ UnitMassMatrix() = UnitMassMatrix{Float64}()
 
 Base.string(::UnitMassMatrix) = "I"
 
-Base.resize!(pc::UnitMassMatrix, θ::AbstractVecOrMat) = nothing
+resize_adaptor!(pc::UnitMassMatrix, size_θ::Tuple) = nothing
 
 reset!(::UnitMassMatrix) = nothing
 
@@ -78,15 +78,16 @@ function get_estimation(nv::NaiveVar)
 end
 
 # Ref： https://github.com/stan-dev/math/blob/develop/stan/math/prim/mat/fun/welford_var_estimator.hpp
-mutable struct WelfordVar{T<:AbstractFloat,E<:AbstractVecOrMat{T}} <: DiagMatrixEstimator{T}
+mutable struct WelfordVar{T<:AbstractFloat,E<:AbstractVecOrMat{T},V<:AbstractVecOrMat{T}} <:
+               DiagMatrixEstimator{T}
     n::Int
     n_min::Int
     μ::E
     M::E
     δ::E    # cache for diff
-    var::E    # cache for variance
-    function WelfordVar(n::Int, n_min::Int, μ::E, M::E, δ::E, var::E) where {E}
-        return new{eltype(E),E}(n, n_min, μ, M, δ, var)
+    var::V    # cache for variance
+    function WelfordVar(n::Int, n_min::Int, μ::E, M::E, δ::E, var::V) where {E,V}
+        return new{eltype(E),E,V}(n, n_min, μ, M, δ, var)
     end
 end
 
@@ -102,20 +103,31 @@ function WelfordVar(sz::Union{Tuple{Int},Tuple{Int,Int}}; kwargs...)
     return WelfordVar{Float64}(sz; kwargs...)
 end
 
-function Base.resize!(wv::WelfordVar, θ::AbstractVecOrMat{T}) where {T<:AbstractFloat}
-    if size(θ) != size(wv.var)
+function resize_adaptor!(wv::WelfordVar{T}, size_θ::Tuple{Int,Int}) where {T<:AbstractFloat}
+    if size_θ != size(wv.var)
         @assert wv.n == 0 "Cannot resize a var estimator when it contains samples."
-        wv.μ = zeros(T, size(θ))
-        wv.M = zeros(T, size(θ))
-        wv.δ = zeros(T, size(θ))
-        wv.var = ones(T, size(θ))
+        wv.μ = zeros(T, size_θ)
+        wv.M = zeros(T, size_θ)
+        wv.δ = zeros(T, size_θ)
+        wv.var = ones(T, size_θ)
+    end
+end
+
+function resize_adaptor!(wv::WelfordVar{T}, size_θ::Tuple{Int}) where {T<:AbstractFloat}
+    length_θ = first(size_θ)
+    if length_θ != size(wv.var, 1)
+        @assert wv.n == 0 "Cannot resize a var estimator when it contains samples."
+        fill!(resize!(wv.μ, length_θ), T(0))
+        fill!(resize!(wv.M, length_θ), T(0))
+        fill!(resize!(wv.δ, length_θ), T(0))
+        fill!(resize!(wv.var, length_θ), T(1))
     end
 end
 
 function reset!(wv::WelfordVar{T}) where {T<:AbstractFloat}
     wv.n = 0
-    wv.μ .= zero(T)
-    wv.M .= zero(T)
+    fill!(wv.μ, zero(T))
+    fill!(wv.M, zero(T))
     return nothing
 end
 
@@ -159,8 +171,6 @@ end
 
 NaiveCov{T}(sz::Tuple{Int}) where {T<:AbstractFloat} = NaiveCov(Vector{Vector{T}}())
 
-NaiveCov(sz::Union{Tuple{Int},Tuple{Int,Int}}; kwargs...) = NaiveCov{Float64}(sz; kwargs...)
-
 Base.push!(nc::NaiveCov, s::AbstractVector) = push!(nc.S, s)
 
 reset!(nc::NaiveCov{T}) where {T} = resize!(nc.S, 0)
@@ -171,13 +181,13 @@ function get_estimation(nc::NaiveCov)
 end
 
 # Ref: https://github.com/stan-dev/math/blob/develop/stan/math/prim/mat/fun/welford_covar_estimator.hpp
-mutable struct WelfordCov{F<:AbstractFloat} <: DenseMatrixEstimator{F}
+mutable struct WelfordCov{F<:AbstractFloat,C<:AbstractMatrix{F}} <: DenseMatrixEstimator{F}
     n::Int
     n_min::Int
     μ::Vector{F}
     M::Matrix{F}
     δ::Vector{F}  # cache for diff
-    cov::Matrix{F}
+    cov::C
 end
 
 Base.show(io::IO, ::WelfordCov) = print(io, "WelfordCov")
@@ -191,20 +201,21 @@ end
 
 WelfordCov(sz::Tuple{Int}; kwargs...) = WelfordCov{Float64}(sz; kwargs...)
 
-function Base.resize!(wc::WelfordCov, θ::AbstractVector{T}) where {T<:AbstractFloat}
-    if length(θ) != size(wc.cov, 1)
+function resize_adaptor!(wc::WelfordCov{T}, size_θ::Tuple{Int}) where {T<:AbstractFloat}
+    length_θ = first(size_θ)
+    if length_θ != size(wc.cov, 1)
         @assert wc.n == 0 "Cannot resize a var estimator when it contains samples."
-        wc.μ = zeros(T, length(θ))
-        wc.δ = zeros(T, length(θ))
-        wc.M = zeros(T, length(θ), length(θ))
-        wc.cov = LinearAlgebra.diagm(0 => ones(T, length(θ)))
+        fill!(resize!(wc.μ, length_θ), T(0))
+        fill!(resize!(wc.δ, length_θ), T(0))
+        wc.M = zeros(T, length_θ, length_θ)
+        wc.cov = LinearAlgebra.diagm(0 => ones(T, length_θ))
     end
 end
 
 function reset!(wc::WelfordCov{T}) where {T<:AbstractFloat}
     wc.n = 0
-    wc.μ .= zero(T)
-    wc.M .= zero(T)
+    fill!(wc.μ, zero(T))
+    fill!(wc.M, zero(T))
     return nothing
 end
 

src/adaptation/massmatrix.jl

@@ -117,7 +117,7 @@ function adapt!(
 
     adapt!(tp.ssa, θ, α)
 
-    resize!(tp.pc, θ) # Resize pre-conditioner if necessary.
+    resize_adaptor!(tp.pc, size(θ)) # Resize pre-conditioner if necessary.
 
     # Ref: https://github.com/stan-dev/stan/blob/develop/src/stan/mcmc/hmc/nuts/adapt_diag_e_nuts.hpp
     if is_in_window(tp)