Use ALL_MODELS rather than DEMO_MODELS (#1173)

* Use ALL_MODELS rather than DEMO_MODELS

* Switch ALL_MODELS for DEMO_MODELS where necessary

* Curb overzealous use of ALL_MODELS

* Add non-BangBang invlink and link for StaticTransformation

* Give demo_lkjchol a non-flat prior PDF

* Mark demo_lkjchol JET test as broken

* Mark an SVI test as broken

* Work around Cholesky comparison bug

* Resolve method ambiguities

* Make a test more robust.

Co-authored-by: Penelope Yong <penelopeysm@gmail.com>

* Fix use of SimpleVarInfo{Dict}

---------

Co-authored-by: Penelope Yong <penelopeysm@gmail.com>

Author: mhauru

docs/src/api.md

@@ -267,6 +267,7 @@ DynamicPPL provides several demo models in the `DynamicPPL.TestUtils` submodule.
 
 ```@docs
 DynamicPPL.TestUtils.DEMO_MODELS
+DynamicPPL.TestUtils.ALL_MODELS
 ```
 
 For every demo model, one can define the true log prior, log likelihood, and log joint probabilities.

src/abstract_varinfo.jl

@@ -837,6 +837,10 @@ end
 function link!!(
     t::StaticTransformation{<:Bijectors.Transform}, vi::AbstractVarInfo, ::Model
 )
+    # TODO(mhauru) This assumes that the user has defined the bijector using the same
+    # variable ordering as what `vi[:]` and `unflatten(vi, x)` use. This is a bad user
+    # interface, and it's also dangerous for any AbstractVarInfo types that may not respect
+    # a particular ordering, such as SimpleVarInfo{Dict}.
     b = inverse(t.bijector)
     x = vi[:]
     y, logjac = with_logabsdet_jacobian(b, x)
@@ -866,7 +870,7 @@ end
 function link(vi::AbstractVarInfo, vns::VarNameTuple, model::Model)
     return link(default_transformation(model, vi), vi, vns, model)
 end
-function link(t::DynamicTransformation, vi::AbstractVarInfo, model::Model)
+function link(t::AbstractTransformation, vi::AbstractVarInfo, model::Model)
     return link!!(t, deepcopy(vi), model)
 end
 
@@ -932,7 +936,7 @@ end
 function invlink(vi::AbstractVarInfo, vns::VarNameTuple, model::Model)
     return invlink(default_transformation(model, vi), vi, vns, model)
 end
-function invlink(t::DynamicTransformation, vi::AbstractVarInfo, model::Model)
+function invlink(t::AbstractTransformation, vi::AbstractVarInfo, model::Model)
     return invlink!!(t, deepcopy(vi), model)
 end
 

src/test_utils/model_interface.jl

@@ -92,7 +92,9 @@ Even though it is recommended to implement this by hand for a particular `Model`
 a default implementation using [`SimpleVarInfo{<:Dict}`](@ref) is provided.
 """
 function varnames(model::Model)
-    return collect(keys(last(DynamicPPL.init!!(model, SimpleVarInfo(Dict())))))
+    result = collect(keys(last(DynamicPPL.init!!(model, SimpleVarInfo(OrderedDict())))))
+    # Concretise the element type.
+    return [x for x in result]
 end
 
 """

src/test_utils/models.jl

@@ -34,6 +34,11 @@ function logprior_true_with_logabsdet_jacobian(
     x_unconstrained, Δlogp = Bijectors.with_logabsdet_jacobian(b_x, x)
     return (m=m, x=x_unconstrained), logprior_true(model, m, x) - Δlogp
 end
+function rand_prior_true(rng::Random.AbstractRNG, ::Model{typeof(demo_dynamic_constraint)})
+    m = rand(rng, Normal())
+    x = rand(rng, truncated(Normal(); lower=m))
+    return (m=m, x=x)
+end
 
 """
     demo_one_variable_multiple_constraints()
@@ -109,12 +114,12 @@ x ~ LKJCholesky(d, 1.0)
 ```
 """
 @model function demo_lkjchol(d::Int=2)
-    x ~ LKJCholesky(d, 1.0)
+    x ~ LKJCholesky(d, 1.5)
     return (x=x,)
 end
 
 function logprior_true(model::Model{typeof(demo_lkjchol)}, x)
-    return logpdf(LKJCholesky(model.args.d, 1.0), x)
+    return logpdf(LKJCholesky(model.args.d, 1.5), x)
 end
 
 function loglikelihood_true(model::Model{typeof(demo_lkjchol)}, x)
@@ -163,6 +168,9 @@ end
 function loglikelihood_true(::Model{typeof(demo_static_transformation)}, s, m)
     return logpdf(Normal(m, sqrt(s)), 1.5) + logpdf(Normal(m, sqrt(s)), 2.0)
 end
+function varnames(::Model{typeof(demo_static_transformation)})
+    return [@varname(s), @varname(m)]
+end
 function logprior_true_with_logabsdet_jacobian(
     model::Model{typeof(demo_static_transformation)}, s, m
 )
@@ -557,22 +565,6 @@ function varnames(model::Model{typeof(demo_assume_matrix_observe_matrix_index)})
     return [@varname(s), @varname(m)]
 end
 
-const DemoModels = Union{
-    Model{typeof(demo_dot_assume_observe)},
-    Model{typeof(demo_assume_index_observe)},
-    Model{typeof(demo_assume_multivariate_observe)},
-    Model{typeof(demo_dot_assume_observe_index)},
-    Model{typeof(demo_assume_dot_observe)},
-    Model{typeof(demo_assume_dot_observe_literal)},
-    Model{typeof(demo_assume_observe_literal)},
-    Model{typeof(demo_assume_multivariate_observe_literal)},
-    Model{typeof(demo_dot_assume_observe_index_literal)},
-    Model{typeof(demo_assume_submodel_observe_index_literal)},
-    Model{typeof(demo_dot_assume_observe_submodel)},
-    Model{typeof(demo_dot_assume_observe_matrix_index)},
-    Model{typeof(demo_assume_matrix_observe_matrix_index)},
-}
-
 const UnivariateAssumeDemoModels = Union{
     Model{typeof(demo_assume_dot_observe)},
     Model{typeof(demo_assume_dot_observe_literal)},
@@ -758,3 +750,14 @@ const DEMO_MODELS = (
     demo_dot_assume_observe_matrix_index(),
     demo_assume_matrix_observe_matrix_index(),
 )
+
+"""
+A tuple of all models defined in DynamicPPL.TestUtils.
+"""
+const ALL_MODELS = (
+    DEMO_MODELS...,
+    demo_dynamic_constraint(),
+    demo_one_variable_multiple_constraints(),
+    demo_lkjchol(),
+    demo_static_transformation(),
+)

src/test_utils/varinfo.jl

@@ -10,7 +10,14 @@ Test that `vi[vn]` corresponds to the correct value in `vals` for every `vn` in
 """
 function test_values(vi::AbstractVarInfo, vals::NamedTuple, vns; compare=isequal, kwargs...)
     for vn in vns
-        @test compare(vi[vn], get(vals, vn); kwargs...)
+        val = get(vals, vn)
+        # TODO(mhauru) Workaround for https://github.com/JuliaLang/LinearAlgebra.jl/pull/1404
+        # Remove once the fix is all Julia versions we support.
+        if val isa Cholesky
+            @test compare(vi[vn].L, val.L; kwargs...)
+        else
+            @test compare(vi[vn], val; kwargs...)
+        end
     end
 end
 

src/threadsafe.jl

@@ -85,12 +85,24 @@ function invlink!!(t::AbstractTransformation, vi::ThreadSafeVarInfo, args...)
     return Accessors.@set vi.varinfo = invlink!!(t, vi.varinfo, args...)
 end
 
-function link(t::AbstractTransformation, vi::ThreadSafeVarInfo, args...)
-    return Accessors.@set vi.varinfo = link(t, vi.varinfo, args...)
+function link(t::AbstractTransformation, vi::ThreadSafeVarInfo, model::Model)
+    return Accessors.@set vi.varinfo = link(t, vi.varinfo, model)
 end
 
-function invlink(t::AbstractTransformation, vi::ThreadSafeVarInfo, args...)
-    return Accessors.@set vi.varinfo = invlink(t, vi.varinfo, args...)
+function invlink(t::AbstractTransformation, vi::ThreadSafeVarInfo, model::Model)
+    return Accessors.@set vi.varinfo = invlink(t, vi.varinfo, model)
+end
+
+function link(
+    t::AbstractTransformation, vi::ThreadSafeVarInfo, vns::VarNameTuple, model::Model
+)
+    return Accessors.@set vi.varinfo = link(t, vi.varinfo, vns, model)
+end
+
+function invlink(
+    t::AbstractTransformation, vi::ThreadSafeVarInfo, vns::VarNameTuple, model::Model
+)
+    return Accessors.@set vi.varinfo = invlink(t, vi.varinfo, vns, model)
 end
 
 # Need to define explicitly for `DynamicTransformation` to avoid method ambiguity.

test/chains.jl

@@ -67,7 +67,7 @@ using Test
 end
 
 @testset "ParamsWithStats from LogDensityFunction" begin
-    @testset "$(m.f)" for m in DynamicPPL.TestUtils.DEMO_MODELS
+    @testset "$(m.f)" for m in DynamicPPL.TestUtils.ALL_MODELS
         unlinked_vi = VarInfo(m)
         @testset "$islinked" for islinked in (false, true)
             vi = if islinked

test/contexts.jl

@@ -179,7 +179,7 @@ Base.IteratorEltype(::Type{<:AbstractContext}) = Base.EltypeUnknown()
             @test new_ctx == FixedContext((b=4,), ConditionContext((a=1,)))
         end
 
-        @testset "evaluation: $(model.f)" for model in DynamicPPL.TestUtils.DEMO_MODELS
+        @testset "evaluation: $(model.f)" for model in DynamicPPL.TestUtils.ALL_MODELS
             prefix_vn = @varname(my_prefix)
             context = DynamicPPL.PrefixContext(prefix_vn, DefaultContext())
             new_model = contextualize(model, context)
@@ -423,7 +423,7 @@ Base.IteratorEltype(::Type{<:AbstractContext}) = Base.EltypeUnknown()
                 DynamicPPL.typed_vector_varinfo(DynamicPPL.typed_varinfo(VarInfo())),
             ),
             ("SVI+NamedTuple", SimpleVarInfo()),
-            ("Svi+Dict", SimpleVarInfo(Dict{VarName,Any}())),
+            ("Svi+Dict", SimpleVarInfo(OrderedDict{VarName,Any}())),
         ]
 
         @model function test_init_model()

test/ext/DynamicPPLJETExt.jl

@@ -61,7 +61,21 @@
     end
 
     @testset "demo models" begin
-        @testset "$(model.f)" for model in DynamicPPL.TestUtils.DEMO_MODELS
+        @testset "$(model.f)" for model in DynamicPPL.TestUtils.ALL_MODELS
+            if model.f === DynamicPPL.TestUtils.demo_lkjchol
+                # TODO(mhauru)
+                # The LKJCholesky model fails with JET. The problem is not with Turing but
+                # with Distributions, and ultimately this in LinearAlgebra:
+                # julia> v = @view rand(2,2)[:,1];
+                #
+                # julia> JET.@report_call norm(v)
+                # ═════ 2 possible errors found ═════
+                # blahblah
+                # The below trivial call to @test is just marking that there's something
+                # broken here.
+                @test false broken = true
+                continue
+            end
             # Use debug logging below.
             varinfo = DynamicPPL.Experimental.determine_suitable_varinfo(model)
             # Check that the inferred varinfo is indeed suitable for evaluation

test/integration/enzyme/main.jl

@@ -1,4 +1,4 @@
-using DynamicPPL.TestUtils: DEMO_MODELS
+using DynamicPPL.TestUtils: ALL_MODELS
 using DynamicPPL.TestUtils.AD: run_ad
 using ADTypes: AutoEnzyme
 using Test: @test, @testset
@@ -17,7 +17,7 @@ ADTYPES = (
 )
 
 @testset "$ad_key" for (ad_key, ad_type) in ADTYPES
-    @testset "$(model.f)" for model in DEMO_MODELS
+    @testset "$(model.f)" for model in ALL_MODELS
         @test run_ad(model, ad_type) isa Any
     end
 end