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)