Merge branch 'mhauru/vnt-for-fastldf' into mhauru/arraylikeblock

Author: mhauru

HISTORY.md

@@ -2,6 +2,19 @@
 
 ## 0.40
 
+## 0.39.4
+
+Removed the internal functions `DynamicPPL.getranges`, `DynamicPPL.vector_getrange`, and `DynamicPPL.vector_getranges` (the new LogDensityFunction construction does exactly the same thing, so this specialised function was not needed).
+
+## 0.39.3
+
+`DynamicPPL.TestUtils.AD.run_ad` now generates much prettier output.
+In particular, when a test fails, it also tells you the tolerances needed to make it pass.
+
+## 0.39.2
+
+`returned(model, parameters...)` now accepts any arguments that can be wrapped in `InitFromParams` (previously it would only accept `NamedTuple`, `AbstractDict{<:VarName}`, or a chain).
+
 ## 0.39.1
 
 `LogDensityFunction` now allows you to call `logdensity_and_gradient(ldf, x)` with `AbstractVector`s `x` that are not plain Vectors (they will be converted internally before calculating the gradient).

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.

ext/DynamicPPLMarginalLogDensitiesExt.jl

@@ -101,12 +101,23 @@ function DynamicPPL.marginalize(
     method::MarginalLogDensities.AbstractMarginalizer=MarginalLogDensities.LaplaceApprox(),
     kwargs...,
 )
-    # Determine the indices for the variables to marginalise out.
-    varindices = reduce(vcat, DynamicPPL.vector_getranges(varinfo, marginalized_varnames))
     # Construct the marginal log-density model.
-    f = DynamicPPL.LogDensityFunction(model, getlogprob, varinfo)
+    ldf = DynamicPPL.LogDensityFunction(model, getlogprob, varinfo)
+    # Determine the indices for the variables to marginalise out.
+    varindices = mapreduce(vcat, marginalized_varnames) do vn
+        if DynamicPPL.getoptic(vn) === identity
+            ldf._iden_varname_ranges[DynamicPPL.getsym(vn)].range
+        else
+            ldf._varname_ranges[vn].range
+        end
+    end
     mld = MarginalLogDensities.MarginalLogDensity(
-        LogDensityFunctionWrapper(f, varinfo), varinfo[:], varindices, (), method; kwargs...
+        LogDensityFunctionWrapper(ldf, varinfo),
+        varinfo[:],
+        varindices,
+        (),
+        method;
+        kwargs...,
     )
     return mld
 end

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/model.jl

@@ -1153,10 +1153,16 @@ end
 function predict end
 
 """
-    returned(model::Model, parameters::NamedTuple)
-    returned(model::Model, parameters::AbstractDict{<:VarName})
+    returned(model::Model, parameters...)
 
-Execute `model` with variables `keys` set to `values` and return the values returned by the `model`.
+Initialise a `model` using the given `parameters` and return the model's return value. The
+parameters must be provided in a format that can be wrapped in an `InitFromParams`, i.e.,
+`InitFromParams(parameters..., nothing)` must be a valid `AbstractInitStrategy` (where
+`nothing` is the fallback strategy to use if parameters are not provided).
+
+As far as DynamicPPL is concerned, `parameters` can be either a singular `NamedTuple` or an
+`AbstractDict{<:VarName}`; however this method is left flexible to allow for other packages
+that wish to extend `InitFromParams`.
 
 # Example
 ```jldoctest
@@ -1177,7 +1183,7 @@ julia> returned(model, Dict{VarName,Float64}(@varname(m) => 2.0))
 (mp1 = 3.0,)
 ```
 """
-function returned(model::Model, parameters::Union{NamedTuple,AbstractDict{<:VarName}})
+function returned(model::Model, parameters...)
     # Note: we can't use `fix(model, parameters)` because
     # https://github.com/TuringLang/DynamicPPL.jl/issues/1097
     return first(
@@ -1186,7 +1192,7 @@ function returned(model::Model, parameters::Union{NamedTuple,AbstractDict{<:VarN
             DynamicPPL.OnlyAccsVarInfo(DynamicPPL.AccumulatorTuple()),
             # Use `nothing` as the fallback to ensure that any missing parameters cause an
             # error
-            InitFromParams(parameters, nothing),
+            InitFromParams(parameters..., nothing),
         ),
     )
 end

src/test_utils/ad.jl

@@ -12,6 +12,7 @@ using DynamicPPL:
     AbstractVarInfo,
     getlogjoint_internal,
     link
+using LinearAlgebra: norm
 using LogDensityProblems: logdensity, logdensity_and_gradient
 using Random: AbstractRNG, default_rng
 using Statistics: median
@@ -78,6 +79,51 @@ struct ADIncorrectException{T<:AbstractFloat} <: Exception
     value_actual::T
     grad_expected::Vector{T}
     grad_actual::Vector{T}
+    atol::T
+    rtol::T
+end
+function Base.showerror(io::IO, e::ADIncorrectException)
+    value_passed = isapprox(e.value_expected, e.value_actual; atol=e.atol, rtol=e.rtol)
+    grad_passed = isapprox(e.grad_expected, e.grad_actual; atol=e.atol, rtol=e.rtol)
+    s = if !value_passed && !grad_passed
+        "value and gradient"
+    elseif !value_passed
+        "value"
+    else
+        "gradient"
+    end
+    println(io, "ADIncorrectException: The AD backend returned an incorrect $s.")
+    println(io, "  Testing was carried out with")
+    println(io, "               atol : $(e.atol)")
+    println(io, "               rtol : $(e.rtol)")
+    # calculate what tolerances would have been needed to pass for value
+    if !value_passed
+        min_atol_needed_to_pass_value = abs(e.value_expected - e.value_actual)
+        min_rtol_needed_to_pass_value =
+            min_atol_needed_to_pass_value / max(abs(e.value_expected), abs(e.value_actual))
+        println(io, "  The value check failed because:")
+        println(io, "     expected value : $(e.value_expected)")
+        println(io, "       actual value : $(e.value_actual)")
+        println(io, "  This value correctness check would have passed if either:")
+        println(io, "               atol ≥ $(min_atol_needed_to_pass_value), or")
+        println(io, "               rtol ≥ $(min_rtol_needed_to_pass_value)")
+    end
+    if !grad_passed
+        norm_expected = norm(e.grad_expected)
+        norm_actual = norm(e.grad_actual)
+        max_norm = max(norm_expected, norm_actual)
+        norm_diff = norm(e.grad_expected - e.grad_actual)
+        min_atol_needed_to_pass_grad = norm_diff
+        min_rtol_needed_to_pass_grad = norm_diff / max_norm
+        # min tolerances needed to pass overall
+        println(io, "  The gradient check failed because:")
+        println(io, "      expected grad : $(e.grad_expected)")
+        println(io, "        actual grad : $(e.grad_actual)")
+        println(io, "  The gradient correctness check would have passed if either:")
+        println(io, "               atol ≥ $(min_atol_needed_to_pass_grad), or")
+        println(io, "               rtol ≥ $(min_rtol_needed_to_pass_grad)")
+    end
+    return nothing
 end
 
 """
@@ -116,11 +162,26 @@ struct ADResult{Tparams<:AbstractFloat,Tresult<:AbstractFloat,Ttol<:AbstractFloa
     "The gradient of logp (calculated using `adtype`)"
     grad_actual::Vector{Tresult}
     "If benchmarking was requested, the time taken by the AD backend to evaluate the gradient
-     of logp"
+    of logp (in seconds)"
     grad_time::Union{Nothing,Tresult}
-    "If benchmarking was requested, the time taken by the AD backend to evaluate logp"
+    "If benchmarking was requested, the time taken by the AD backend to evaluate logp (in
+    seconds)"
     primal_time::Union{Nothing,Tresult}
 end
+function Base.show(io::IO, ::MIME"text/plain", result::ADResult)
+    printstyled(io, "ADResult\n"; bold=true)
+    println(io, "  ├ model          : $(result.model.f)")
+    println(io, "  ├ adtype         : $(result.adtype)")
+    println(io, "  ├ value_actual   : $(result.value_actual)")
+    println(io, "  ├ value_expected : $(result.value_expected)")
+    println(io, "  ├ grad_actual    : $(result.grad_actual)")
+    println(io, "  ├ grad_expected  : $(result.grad_expected)")
+    if result.grad_time !== nothing && result.primal_time !== nothing
+        println(io, "  ├ grad_time      : $(result.grad_time) s")
+        println(io, "  ├ primal_time    : $(result.primal_time) s")
+    end
+    return println(io, "  └ params         : $(result.params)")
+end
 
 """
     run_ad(
@@ -230,6 +291,14 @@ Everything else is optional, and can be categorised into several groups:
    we cannot know the magnitude of logp and its gradient a priori. The `atol`
    value is supplied to handle the case where gradients are equal to zero.
 
+1. _Whether to benchmark._
+
+   By default, benchmarking is disabled. To enable it, set `benchmark=true`.
+   When enabled, the time taken to evaluate logp as well as its gradient is
+   measured using Chairmarks.jl, and the `ADResult` object returned will
+   contain `grad_time` and `primal_time` fields with the median times (in
+   seconds).
+
 1. _Whether to output extra logging information._
 
    By default, this function prints messages when it runs. To silence it, set
@@ -297,7 +366,7 @@ function run_ad(
         end
         # Perform testing
         verbose && println("     expected : $((value_true, grad_true))")
-        exc() = throw(ADIncorrectException(value, value_true, grad, grad_true))
+        exc() = throw(ADIncorrectException(value, value_true, grad, grad_true, atol, rtol))
         isapprox(value, value_true; atol=atol, rtol=rtol) || exc()
         isapprox(grad, grad_true; atol=atol, rtol=rtol) || exc()
     end
@@ -306,14 +375,18 @@ function run_ad(
     grad_time, primal_time = if benchmark
         logdensity(ldf, params)  # Warm-up
         primal_benchmark = @be logdensity($ldf, $params)
+        print("   evaluation : ")
+        show(stdout, MIME("text/plain"), median(primal_benchmark))
+        println()
         logdensity_and_gradient(ldf, params)  # Warm-up
         grad_benchmark = @be logdensity_and_gradient($ldf, $params)
+        print("     gradient : ")
+        show(stdout, MIME("text/plain"), median(grad_benchmark))
+        println()
         median_primal = median(primal_benchmark).time
         median_grad = median(grad_benchmark).time
         r(f) = round(f; sigdigits=4)
-        verbose && println(
-            "grad / primal : $(r(median_grad))/$(r(median_primal)) = $(r(median_grad / median_primal))",
-        )
+        verbose && println("  grad / eval : $(r(median_grad / median_primal))")
         (median_grad, median_primal)
     else
         nothing, nothing

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.
@@ -155,10 +167,6 @@ function BangBang.setindex!!(vi::ThreadSafeVarInfo, vals, vns::AbstractVector{<:
 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
 
 isempty(vi::ThreadSafeVarInfo) = isempty(vi.varinfo)
 function BangBang.empty!!(vi::ThreadSafeVarInfo)