Add Aqua tests (#775)

* Add Aqua tests

* Fix logpdf(::NamedDist) method ambiguity

* Fix SimpleVarInfo method ambiguity

* Fix VarInfo method ambiguity

* Add InteractiveUtils compat entry

See: https://discourse.julialang.org/t/psa-compat-requirements-in-the-general-registry-are-changing/104958

* Add Random.AbstractRNG type annotation

* Remove unneeded getsym method

* Fix (newly introduced 😅) ConditionContext method ambiguity

* Fix unwrap_right_left_vns method ambiguity

* KernelAbstractions is a weakdep not a dep

* Fix StaticTransformation / ThreadSafeVarInfo link/invlink ambiguity

* Fix more RNGs

* Don't run Aqua tests on CI min versions

* Fix ternary in GitHub Actions expression

Author: penelopeysm

.github/workflows/CI.yml

@@ -70,6 +70,8 @@ jobs:
         env:
           GROUP: ${{ matrix.test_group }}
           JULIA_NUM_THREADS: ${{ matrix.runner.num_threads }}
+          # Only run Aqua tests on latest version
+          AQUA: ${{ matrix.runner.version == '1' && 'true' || 'false' }}
 
       - uses: julia-actions/julia-processcoverage@v1
 

Project.toml

@@ -16,7 +16,6 @@ DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
-KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
@@ -30,6 +29,7 @@ ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
+KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 
@@ -56,6 +56,7 @@ Distributions = "0.25"
 DocStringExtensions = "0.9"
 EnzymeCore = "0.6 - 0.8"
 ForwardDiff = "0.10.12"
+InteractiveUtils = "1"
 JET = "0.9"
 KernelAbstractions = "0.9.33"
 LinearAlgebra = "1.6"

benchmarks/src/DynamicPPLBenchmarks.jl

@@ -9,6 +9,7 @@ using LogDensityProblems: LogDensityProblems
 using ForwardDiff: ForwardDiff
 using Mooncake: Mooncake
 using ReverseDiff: ReverseDiff
+using StableRNGs: StableRNG
 
 include("./Models.jl")
 using .Models: Models
@@ -61,18 +62,20 @@ The AD backend should be specified as a Symbol (e.g. `:forwarddiff`, `:reversedi
 `islinked` determines whether to link the VarInfo for evaluation.
 """
 function make_suite(model, varinfo_choice::Symbol, adbackend::Symbol, islinked::Bool)
+    rng = StableRNG(23)
+
     suite = BenchmarkGroup()
 
     vi = if varinfo_choice == :untyped
         vi = VarInfo()
-        model(vi)
+        model(rng, vi)
         vi
     elseif varinfo_choice == :typed
-        VarInfo(model)
+        VarInfo(rng, model)
     elseif varinfo_choice == :simple_namedtuple
-        SimpleVarInfo{Float64}(model())
+        SimpleVarInfo{Float64}(model(rng))
     elseif varinfo_choice == :simple_dict
-        retvals = model()
+        retvals = model(rng)
         vns = [VarName{k}() for k in keys(retvals)]
         SimpleVarInfo{Float64}(Dict(zip(vns, values(retvals))))
     else

src/compiler.jl

@@ -250,7 +250,10 @@ x[1][3]
 ```
 """
 unwrap_right_left_vns(right, left, vns) = right, left, vns
-function unwrap_right_left_vns(right::NamedDist, left, vns)
+function unwrap_right_left_vns(right::NamedDist, left::AbstractArray, ::VarName)
+    return unwrap_right_left_vns(right.dist, left, right.name)
+end
+function unwrap_right_left_vns(right::NamedDist, left::AbstractMatrix, ::VarName)
     return unwrap_right_left_vns(right.dist, left, right.name)
 end
 function unwrap_right_left_vns(

src/context_implementations.jl

@@ -195,7 +195,7 @@ function tilde_observe!!(context, right, left, vi)
     return left, acclogp_observe!!(context, vi, logp)
 end
 
-function assume(rng, spl::Sampler, dist)
+function assume(rng::Random.AbstractRNG, spl::Sampler, dist)
     return error("DynamicPPL.assume: unmanaged inference algorithm: $(typeof(spl))")
 end
 

src/contexts.jl

@@ -335,6 +335,8 @@ function ConditionContext(values::Union{NamedTuple,AbstractDict})
 end
 # Optimisation when there are no values to condition on
 ConditionContext(::NamedTuple{()}, context::AbstractContext) = context
+# Same as above, and avoids method ambiguity with below
+ConditionContext(::NamedTuple{()}, context::NamedConditionContext) = context
 # Collapse consecutive levels of `ConditionContext`. Note that this overrides
 # values inside the child context, thus giving precedence to the outermost
 # `ConditionContext`.

src/distribution_wrappers.jl

@@ -17,6 +17,10 @@ Base.length(dist::NamedDist) = Base.length(dist.dist)
 Base.size(dist::NamedDist) = Base.size(dist.dist)
 
 Distributions.logpdf(dist::NamedDist, x::Real) = Distributions.logpdf(dist.dist, x)
+function Distributions.logpdf(dist::NamedDist, x::AbstractArray{<:Real,0})
+    # extract the singleton value from 0-dimensional array
+    return Distributions.logpdf(dist.dist, first(x))
+end
 function Distributions.logpdf(dist::NamedDist, x::AbstractArray{<:Real})
     return Distributions.logpdf(dist.dist, x)
 end

src/simple_varinfo.jl

@@ -232,8 +232,14 @@ function SimpleVarInfo(; kwargs...)
 end
 
 # Constructor from `Model`.
-SimpleVarInfo(model::Model, args...) = SimpleVarInfo{Float64}(model, args...)
-function SimpleVarInfo{T}(model::Model, args...) where {T<:Real}
+function SimpleVarInfo(
+    model::Model, args::Union{AbstractVarInfo,AbstractSampler,AbstractContext}...
+)
+    return SimpleVarInfo{Float64}(model, args...)
+end
+function SimpleVarInfo{T}(
+    model::Model, args::Union{AbstractVarInfo,AbstractSampler,AbstractContext}...
+) where {T<:Real}
     return last(evaluate!!(model, SimpleVarInfo{T}(), args...))
 end
 

src/threadsafe.jl

@@ -115,6 +115,19 @@ function invlink(t::DynamicTransformation, vi::ThreadSafeVarInfo, model::Model)
     return invlink!!(t, deepcopy(vi), model)
 end
 
+# These two StaticTransformation methods needed to resolve ambiguities
+function link!!(
+    t::StaticTransformation{<:Bijectors.Transform}, vi::ThreadSafeVarInfo, model::Model
+)
+    return Accessors.@set vi.varinfo = link!!(t, vi.varinfo, model)
+end
+
+function invlink!!(
+    t::StaticTransformation{<:Bijectors.Transform}, vi::ThreadSafeVarInfo, model::Model
+)
+    return Accessors.@set vi.varinfo = invlink!!(t, vi.varinfo, model)
+end
+
 function maybe_invlink_before_eval!!(vi::ThreadSafeVarInfo, model::Model)
     # Defer to the wrapped `AbstractVarInfo` object.
     # NOTE: When computing `getlogp` for `ThreadSafeVarInfo` we do include the

src/varinfo.jl

@@ -200,7 +200,11 @@ function VarInfo(
 )
     return typed_varinfo(model, SamplingContext(rng, sampler, context), metadata)
 end
-VarInfo(model::Model, args...) = VarInfo(Random.default_rng(), model, args...)
+function VarInfo(
+    model::Model, args::Union{AbstractSampler,AbstractContext,Metadata,VarNamedVector}...
+)
+    return VarInfo(Random.default_rng(), model, args...)
+end
 
 """
     vector_length(varinfo::VarInfo)

src/varname.jl

@@ -41,6 +41,3 @@ Possibly existing indices of `varname` are neglected.
 ) where {s,missings,_F,_a,_T}
     return s in missings
 end
-
-# HACK: Type-piracy. Is this really the way to go?
-AbstractPPL.getsym(::AbstractVector{<:VarName{sym}}) where {sym} = sym

test/Aqua.jl

@@ -0,0 +1,8 @@
+module AquaTests
+
+using Aqua: Aqua
+using DynamicPPL
+
+Aqua.test_all(DynamicPPL)
+
+end

test/Project.toml

@@ -3,6 +3,7 @@ ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
 AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf"
 Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
+Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 Bijectors = "76274a88-744f-5084-9051-94815aaf08c4"
 Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"

test/runtests.jl

@@ -35,6 +35,7 @@ using OrderedCollections: OrderedSet
 using DynamicPPL: getargs_dottilde, getargs_tilde
 
 const GROUP = get(ENV, "GROUP", "All")
+const AQUA = get(ENV, "AQUA", "true") == "true"
 Random.seed!(100)
 
 include("test_util.jl")
@@ -44,6 +45,9 @@ include("test_util.jl")
     # groups are chosen to make both groups take roughly the same amount of
     # time, but beyond that there is no particular reason for the split.
     if GROUP == "All" || GROUP == "Group1"
+        if AQUA
+            include("Aqua.jl")
+        end
         include("utils.jl")
         include("compiler.jl")
         include("varnamedvector.jl")