Merge branch 'master' into torfjelde/determine-varinfo

Author: Tor Fjelde

Project.toml

@@ -1,6 +1,6 @@
 name = "DynamicPPL"
 uuid = "366bfd00-2699-11ea-058f-f148b4cae6d8"
-version = "0.31.1"
+version = "0.31.4"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -31,7 +31,7 @@ EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
-ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
+Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
 
 [extensions]
@@ -40,7 +40,7 @@ DynamicPPLEnzymeCoreExt = ["EnzymeCore"]
 DynamicPPLForwardDiffExt = ["ForwardDiff"]
 DynamicPPLJETExt = ["JET"]
 DynamicPPLMCMCChainsExt = ["MCMCChains"]
-DynamicPPLReverseDiffExt = ["ReverseDiff"]
+DynamicPPLMooncakeExt = ["Mooncake"]
 DynamicPPLZygoteRulesExt = ["ZygoteRules"]
 
 [compat]
@@ -63,19 +63,10 @@ LogDensityProblems = "2"
 LogDensityProblemsAD = "1.7.0"
 MCMCChains = "6"
 MacroTools = "0.5.6"
+Mooncake = "0.4.59"
 OrderedCollections = "1"
 Random = "1.6"
 Requires = "1"
-ReverseDiff = "1"
 Test = "1.6"
 ZygoteRules = "0.2"
 julia = "1.10"
-
-[extras]
-ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
-EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
-ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
-JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
-MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
-ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
-ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"

ext/DynamicPPLMooncakeExt.jl

@@ -0,0 +1,9 @@
+module DynamicPPLMooncakeExt
+
+using DynamicPPL: DynamicPPL, istrans
+using Mooncake: Mooncake
+
+# This is purely an optimisation.
+Mooncake.@zero_adjoint Mooncake.DefaultCtx Tuple{typeof(istrans),Vararg}
+
+end # module

ext/DynamicPPLReverseDiffExt.jl

@@ -144,3 +144,19 @@ function LogDensityProblems.capabilities(::Type{<:LogDensityFunction})
 end
 # TODO: should we instead implement and call on `length(f.varinfo)` (at least in the cases where no sampler is involved)?
 LogDensityProblems.dimension(f::LogDensityFunction) = length(getparams(f))
+
+# This is important for performance -- one needs to provide `ADGradient` with a vector of
+# parameters, or DifferentiationInterface will not have sufficient information to e.g.
+# compile a rule for Mooncake (because it won't know the type of the input), or pre-allocate
+# a tape when using ReverseDiff.jl.
+function _make_ad_gradient(ad::ADTypes.AbstractADType, ℓ::LogDensityFunction)
+    x = map(identity, getparams(ℓ)) # ensure we concretise the elements of the params
+    return LogDensityProblemsAD.ADgradient(ad, ℓ; x)
+end
+
+function LogDensityProblemsAD.ADgradient(ad::ADTypes.AutoMooncake, f::LogDensityFunction)
+    return _make_ad_gradient(ad, f)
+end
+function LogDensityProblemsAD.ADgradient(ad::ADTypes.AutoReverseDiff, f::LogDensityFunction)
+    return _make_ad_gradient(ad, f)
+end

src/logdensityfunction.jl

@@ -178,6 +178,12 @@ function BangBang.setindex!!(vi::ThreadSafeVarInfo, vals, vns::AbstractVector{<:
     return Accessors.@set vi.varinfo = BangBang.setindex!!(vi.varinfo, vals, vns)
 end
 
+vector_length(vi::ThreadSafeVarInfo) = vector_length(vi.varinfo)
+vector_getrange(vi::ThreadSafeVarInfo, vn::VarName) = vector_getrange(vi.varinfo, vn)
+function vector_getranges(vi::ThreadSafeVarInfo, vns::Vector{<:VarName})
+    return vector_getranges(vi.varinfo, vns)
+end
+
 function set_retained_vns_del_by_spl!(vi::ThreadSafeVarInfo, spl::Sampler)
     return set_retained_vns_del_by_spl!(vi.varinfo, spl)
 end

src/threadsafe.jl

@@ -208,6 +208,15 @@ function VarInfo(
 end
 VarInfo(model::Model, args...) = VarInfo(Random.default_rng(), model, args...)
 
+"""
+    vector_length(varinfo::VarInfo)
+
+Return the length of the vector representation of `varinfo`.
+"""
+vector_length(varinfo::VarInfo) = length(varinfo.metadata)
+vector_length(varinfo::TypedVarInfo) = sum(length, varinfo.metadata)
+vector_length(md::Metadata) = sum(length, md.ranges)
+
 unflatten(vi::VarInfo, x::AbstractVector) = unflatten(vi, SampleFromPrior(), x)
 
 # TODO: deprecate.
@@ -632,7 +641,72 @@ setrange!(md::Metadata, vn::VarName, range) = md.ranges[getidx(md, vn)] = range
 Return the indices of `vns` in the metadata of `vi` corresponding to `vn`.
 """
 function getranges(vi::VarInfo, vns::Vector{<:VarName})
-    return mapreduce(vn -> getrange(vi, vn), vcat, vns; init=Int[])
+    return map(Base.Fix1(getrange, vi), vns)
+end
+
+"""
+    vector_getrange(varinfo::VarInfo, varname::VarName)
+
+Return the range corresponding to `varname` in the vector representation of `varinfo`.
+"""
+vector_getrange(vi::VarInfo, vn::VarName) = getrange(vi.metadata, vn)
+function vector_getrange(vi::TypedVarInfo, vn::VarName)
+    offset = 0
+    for md in values(vi.metadata)
+        # First, we need to check if `vn` is in `md`.
+        # In this case, we can just return the corresponding range + offset.
+        haskey(md, vn) && return getrange(md, vn) .+ offset
+        # Otherwise, we need to get the cumulative length of the ranges in `md`
+        # and add it to the offset.
+        offset += sum(length, md.ranges)
+    end
+    # If we reach this point, `vn` is not in `vi.metadata`.
+    throw(KeyError(vn))
+end
+
+"""
+    vector_getranges(varinfo::VarInfo, varnames::Vector{<:VarName})
+
+Return the range corresponding to `varname` in the vector representation of `varinfo`.
+"""
+function vector_getranges(varinfo::VarInfo, varname::Vector{<:VarName})
+    return map(Base.Fix1(vector_getrange, varinfo), varname)
+end
+# Specialized version for `TypedVarInfo`.
+function vector_getranges(varinfo::TypedVarInfo, vns::Vector{<:VarName})
+    # TODO: Does it help if we _don't_ convert to a vector here?
+    metadatas = collect(values(varinfo.metadata))
+    # Extract the offsets.
+    offsets = cumsum(map(vector_length, metadatas))
+    # Extract the ranges from each metadata.
+    ranges = Vector{UnitRange{Int}}(undef, length(vns))
+    # Need to keep track of which ones we've seen.
+    not_seen = fill(true, length(vns))
+    for (i, metadata) in enumerate(metadatas)
+        vns_metadata = filter(Base.Fix1(haskey, metadata), vns)
+        # If none of the variables exist in the metadata, we return an empty array.
+        isempty(vns_metadata) && continue
+        # Otherwise, we extract the ranges.
+        offset = i == 1 ? 0 : offsets[i - 1]
+        for vn in vns_metadata
+            r_vn = getrange(metadata, vn)
+            # Get the index, so we return in the same order as `vns`.
+            # NOTE: There might be duplicates in `vns`, so we need to handle that.
+            indices = findall(==(vn), vns)
+            for idx in indices
+                not_seen[idx] = false
+                ranges[idx] = r_vn .+ offset
+            end
+        end
+    end
+    # Raise key error if any of the variables were not found.
+    if any(not_seen)
+        inds = findall(not_seen)
+        # Just use a `convert` to get the same type as the input; don't want to confuse by overly
+        # specilizing the types in the error message.
+        throw(KeyError(convert(typeof(vns), vns[inds])))
+    end
+    return ranges
 end
 
 """
@@ -1320,13 +1394,13 @@ end
 
 function _inner_transform!(md::Metadata, vi::VarInfo, vn::VarName, f)
     # TODO: Use inplace versions to avoid allocations
-    yvec, logjac = with_logabsdet_jacobian(f, getindex_internal(vi, vn))
+    yvec, logjac = with_logabsdet_jacobian(f, getindex_internal(md, vn))
     # Determine the new range.
-    start = first(getrange(vi, vn))
+    start = first(getrange(md, vn))
     # NOTE: `length(yvec)` should never be longer than `getrange(vi, vn)`.
-    setrange!(vi, vn, start:(start + length(yvec) - 1))
+    setrange!(md, vn, start:(start + length(yvec) - 1))
     # Set the new value.
-    setval!(vi, yvec, vn)
+    setval!(md, yvec, vn)
     acclogp!!(vi, -logjac)
     return vi
 end

src/varinfo.jl

@@ -1036,6 +1036,8 @@ function replace_raw_storage(vnv::VarNamedVector, ::Val{space}, vals) where {spa
     return replace_raw_storage(vnv, vals)
 end
 
+vector_length(vnv::VarNamedVector) = length(vnv.vals) - num_inactive(vnv)
+
 """
     unflatten(vnv::VarNamedVector, vals::AbstractVector)
 

src/varnamedvector.jl

@@ -6,6 +6,7 @@ Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
 Bijectors = "76274a88-744f-5084-9051-94815aaf08c4"
 Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
+DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 DistributionsAD = "ced4e74d-a319-5a8a-b0ac-84af2272839c"
@@ -18,6 +19,7 @@ LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
 LogDensityProblemsAD = "996a588d-648d-4e1f-a8f0-a84b347e47b1"
 MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
+Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
@@ -35,6 +37,7 @@ Accessors = "0.1"
 Bijectors = "0.15.1"
 Combinatorics = "1"
 Compat = "4.3.0"
+DifferentiationInterface = "0.6"
 Distributions = "0.25"
 DistributionsAD = "0.6.3"
 Documenter = "1"
@@ -44,6 +47,7 @@ LogDensityProblems = "2"
 LogDensityProblemsAD = "1.7.0"
 MCMCChains = "6.0.4"
 MacroTools = "0.5.6"
+Mooncake = "0.4.59"
 ReverseDiff = "1"
 StableRNGs = "1"
 Tracker = "0.2.23"

test/Project.toml

@@ -1,4 +1,4 @@
-@testset "AD: ForwardDiff and ReverseDiff" begin
+@testset "AD: ForwardDiff, ReverseDiff, and Mooncake" begin
     @testset "$(m.f)" for m in DynamicPPL.TestUtils.DEMO_MODELS
         f = DynamicPPL.LogDensityFunction(m)
         rand_param_values = DynamicPPL.TestUtils.rand_prior_true(m)
@@ -17,11 +17,20 @@
             θ = convert(Vector{Float64}, varinfo[:])
             logp, ref_grad = LogDensityProblems.logdensity_and_gradient(ad_forwarddiff_f, θ)
 
-            @testset "ReverseDiff with compile=$compile" for compile in (false, true)
-                adtype = ADTypes.AutoReverseDiff(; compile=compile)
-                ad_f = LogDensityProblemsAD.ADgradient(adtype, f)
-                _, grad = LogDensityProblems.logdensity_and_gradient(ad_f, θ)
-                @test grad ≈ ref_grad
+            @testset "$adtype" for adtype in [
+                ADTypes.AutoReverseDiff(; compile=false),
+                ADTypes.AutoReverseDiff(; compile=true),
+                ADTypes.AutoMooncake(; config=nothing),
+            ]
+                # Mooncake can't currently handle something that is going on in
+                # SimpleVarInfo{<:VarNamedVector}. Disable all SimpleVarInfo tests for now.
+                if adtype isa ADTypes.AutoMooncake && varinfo isa DynamicPPL.SimpleVarInfo
+                    @test_broken 1 == 0
+                else
+                    ad_f = LogDensityProblemsAD.ADgradient(adtype, f)
+                    _, grad = LogDensityProblems.logdensity_and_gradient(ad_f, θ)
+                    @test grad ≈ ref_grad
+                end
             end
         end
     end

test/ad.jl

@@ -0,0 +1,5 @@
+@testset "DynamicPPLMooncakeExt" begin
+    Mooncake.TestUtils.test_rule(
+        StableRNG(123456), istrans, VarInfo(); unsafe_perturb=true, interface_only=true
+    )
+end