diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 3a992da5..60609e16 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -52,3 +52,23 @@ jobs: files: lcov.info token: ${{ secrets.CODECOV_TOKEN }} fail_ci_if_error: true + + ext: + name: Ext (logdensityproblems, ${{ matrix.version }}) + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + version: + - '1' + - 'min' + steps: + - uses: actions/checkout@v6 + - uses: julia-actions/setup-julia@v2 + with: + version: ${{ matrix.version }} + - uses: julia-actions/cache@v3 + - uses: julia-actions/julia-buildpkg@v1 + - run: julia --project=. test/run_extras.jl + env: + LABEL: ext/logdensityproblems diff --git a/Project.toml b/Project.toml index 00e48d00..3678edc6 100644 --- a/Project.toml +++ b/Project.toml @@ -3,9 +3,10 @@ uuid = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf" keywords = ["probablistic programming"] license = "MIT" desc = "Common interfaces for probabilistic programming" -version = "0.14.2" +version = "0.14.3" [deps] +ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001" Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697" BangBang = "198e06fe-97b7-11e9-32a5-e1d131e6ad66" @@ -19,11 +20,15 @@ StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" [weakdeps] Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" +LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c" [extensions] AbstractPPLDistributionsExt = ["Distributions", "LinearAlgebra"] +AbstractPPLLogDensityProblemsExt = ["LogDensityProblems"] [compat] +ADTypes = "1" +LogDensityProblems = "2" AbstractMCMC = "2, 3, 4, 5" Accessors = "0.1" BangBang = "0.4" diff --git a/docs/Project.toml b/docs/Project.toml index 9aed942a..1ef0eca3 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -1,8 +1,10 @@ [deps] +ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" +ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" [sources] -AbstractPPL = {path = "../"} +AbstractPPL = {path = ".."} diff --git a/docs/make.jl b/docs/make.jl index abab3f69..ead3820f 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -9,7 +9,7 @@ DocMeta.setdocmeta!(AbstractPPL, :DocTestSetup, :(using AbstractPPL); recursive= makedocs(; sitename="AbstractPPL", modules=[AbstractPPL, Base.get_extension(AbstractPPL, :AbstractPPLDistributionsExt)], - pages=["index.md", "varname.md", "pplapi.md", "interface.md"], + pages=["index.md", "varname.md", "pplapi.md", "evaluators.md", "interface.md"], checkdocs=:exports, doctest=false, ) diff --git a/docs/src/evaluators.md b/docs/src/evaluators.md new file mode 100644 index 00000000..4f1e1512 --- /dev/null +++ b/docs/src/evaluators.md @@ -0,0 +1,163 @@ +# Evaluator preparation and AD + +AbstractPPL provides a small interface for preparing callables and asking a +prepared evaluator for values and derivatives. `prepare` binds a callable to a +sample input that establishes the expected input shape and type; +`value_and_gradient!!` and `value_and_jacobian!!` then return the value and +derivative together. + +The `!!` suffix signals that the returned gradient or Jacobian **may alias +internal cache buffers** of the prepared evaluator. The next call to +`value_and_gradient!!` (or `value_and_jacobian!!`) may overwrite that buffer +in place, so a previously-returned reference will silently change. Copy +before holding on to a result: + +```julia +val, grad = value_and_gradient!!(prepared, x1) +saved = copy(grad) # safe to keep +val2, grad2 = value_and_gradient!!(prepared, x2) +# `grad` may now reflect `x2`; `saved` still reflects `x1` +``` + +Backends that always allocate fresh output (e.g. `ForwardDiff.gradient`) do +not actually alias, but consumers should not rely on that — write to the +contract, not the implementation. + +## Quick start + +```@example ad +using AbstractPPL +using AbstractPPL: prepare, value_and_gradient!! +using AbstractPPL.Evaluators: Prepared, VectorEvaluator, NamedTupleEvaluator +using ADTypes: AutoForwardDiff +using ForwardDiff: ForwardDiff + +function AbstractPPL.prepare(adtype::AutoForwardDiff, f, x::AbstractVector{<:Real}) + return Prepared(adtype, VectorEvaluator(f, length(x))) +end + +function AbstractPPL.value_and_gradient!!( + p::Prepared{<:AutoForwardDiff}, x::AbstractVector{<:Real} +) + return (p(x), ForwardDiff.gradient(p.evaluator.f, x)) +end + +mvnormal_logp(x) = -0.5 * sum(abs2, x) # standard normal log density (up to constant) +prepared = prepare(AutoForwardDiff(), mvnormal_logp, zeros(3)) +value_and_gradient!!(prepared, [1.0, 2.0, 3.0]) +``` + +## Two input styles + +### Vector inputs + +When the callable accepts a flat vector, pass a sample vector whose length +matches the expected input: + +```@example ad +prepared([1.0, 2.0, 3.0]) +``` + +For vector-valued callables, use `value_and_jacobian!!`. The returned Jacobian +has shape `(length(value), length(x))`. The same backend extension that +defines `value_and_gradient!!` typically also defines `value_and_jacobian!!` +on the same `Prepared` type — they are separate generic functions, so the +two methods coexist without conflict and the caller picks whichever applies +to their function: + +```@example ad +using AbstractPPL: value_and_jacobian!! + +function AbstractPPL.value_and_jacobian!!( + p::Prepared{<:AutoForwardDiff}, x::AbstractVector{<:Real} +) + return (p(x), ForwardDiff.jacobian(p.evaluator.f, x)) +end + +vecfun(x) = [x[1] * x[2], x[2] + x[3]] +prepared_vec = prepare(AutoForwardDiff(), vecfun, zeros(3)) +value_and_jacobian!!(prepared_vec, [2.0, 3.0, 4.0]) +``` + +### NamedTuple inputs + +When the callable accepts a `NamedTuple`, pass a sample `NamedTuple` whose +field names and value types match the expected input. The prototype's leaves +must be `Real`, `Complex`, `AbstractArray` (recursively), `Tuple`, or +`NamedTuple`. An extension can define a `prepare` overload that wraps the +function in a `NamedTupleEvaluator`: + +```@example ad +function AbstractPPL.prepare(adtype::AutoForwardDiff, f, values::NamedTuple) + return Prepared(adtype, NamedTupleEvaluator(f, values)) +end + +ntfun(v::NamedTuple) = v.a^2 + sum(abs2, v.b) +prepared_nt = prepare(AutoForwardDiff(), ntfun, (a=0.0, b=zeros(2))) +prepared_nt((a=1.0, b=[2.0, 3.0])) +``` + +## AD backends + +Automatic differentiation packages extend the interface by implementing +`value_and_gradient!!` and `value_and_jacobian!!` for specific cache types +stored in `prepared.cache`: + +```julia +prepared = prepare(adtype, problem, prototype) # returns Prepared{AD,E,Cache} +value_and_gradient!!(prepared, x) # may return aliased cache buffer +value_and_jacobian!!(prepared, x) +``` + +`Prepared` has three fields: `adtype`, `evaluator` (the user-facing callable), +and `cache` (backend-specific pre-allocated state such as ForwardDiff configs or +Mooncake tapes). Backend extensions dispatch on the cache type: + +```julia +function AbstractPPL.prepare( + adtype::MyADType, problem, x::AbstractVector{<:Real}; check_dims::Bool=true +) + f = ... # extract callable from problem + cache = MyCache(f, x) + return Prepared(adtype, VectorEvaluator{check_dims}(f, length(x)), cache) +end + +function AbstractPPL.value_and_gradient!!( + p::Prepared{<:AbstractADType,<:VectorEvaluator,<:MyCache}, x::AbstractVector{<:Real} +) + # use p.cache to avoid allocations + return ... +end +``` + +Pass `check_dims=false` in your `prepare` implementation to construct a +`VectorEvaluator{false}`, which skips the per-call length check. This is an +opt-in trust mode — the caller takes responsibility for `length(x)`. The +typical use is inside a backend's `value_and_gradient!!`, where the AD +library invokes the inner callable many times with same-length dual arrays +derived from a single user-supplied `x`; re-validating on each invocation +would be redundant work in the hot path. + +## Without an AD backend + +The two-argument form `prepare(problem, x)` is available without any AD +package. By default it wraps `problem` in a `VectorEvaluator{check_dims}` +(or `NamedTupleEvaluator{check_dims}` for the `NamedTuple` form), giving you +a callable that runs the per-call shape check before forwarding to +`problem`. Downstream code that only needs primal evaluation (e.g. +log-density only, no gradient) can call `prepare(...)` uniformly without +knowing whether an AD backend is loaded: + +```@example ad +sumsimple(x) = sum(x) +p = prepare(sumsimple, zeros(3)) # `VectorEvaluator{true}(sumsimple, 3)` +p([1.0, 2.0, 3.0]) +``` + +## API reference + +```@docs +AbstractPPL.prepare +AbstractPPL.value_and_gradient!! +AbstractPPL.value_and_jacobian!! +``` diff --git a/docs/src/pplapi.md b/docs/src/pplapi.md index 492ab294..20a4fcf9 100644 --- a/docs/src/pplapi.md +++ b/docs/src/pplapi.md @@ -18,3 +18,7 @@ evaluate!! ```@docs AbstractModelTrace ``` + +## Evaluators interface + +See [Evaluator preparation and AD](@ref) for a full guide and API reference. diff --git a/ext/AbstractPPLLogDensityProblemsExt.jl b/ext/AbstractPPLLogDensityProblemsExt.jl new file mode 100644 index 00000000..8f00ebb7 --- /dev/null +++ b/ext/AbstractPPLLogDensityProblemsExt.jl @@ -0,0 +1,38 @@ +module AbstractPPLLogDensityProblemsExt + +using AbstractPPL: AbstractPPL +using AbstractPPL.Evaluators: Prepared, VectorEvaluator +using LogDensityProblems: LogDensityProblems + +# LDP integration is restricted to vector-input evaluators; `NamedTupleEvaluator` +# does not satisfy LDP's vector-input contract. Scalar output is a runtime +# contract the user must satisfy. + +LogDensityProblems.logdensity(p::Prepared{<:Any,<:VectorEvaluator}, x) = p(x) +LogDensityProblems.logdensity(e::VectorEvaluator, x) = e(x) + +function LogDensityProblems.dimension(p::Prepared{<:Any,<:VectorEvaluator}) + return LogDensityProblems.dimension(p.evaluator) +end +LogDensityProblems.dimension(e::VectorEvaluator) = e.dim + +# Order 0 by default. AD-backend extensions (DifferentiationInterface, +# ForwardDiff, Mooncake, etc.) must overload this for their cache type to +# advertise `LogDensityOrder{1}` — without that overload, +# `logdensity_and_gradient` will hit the `value_and_gradient!!` stub and fail. +# LDP defines `capabilities(x) = capabilities(typeof(x))`, so the type method +# alone covers both call shapes. +function LogDensityProblems.capabilities(::Type{<:Prepared{<:Any,<:VectorEvaluator}}) + return LogDensityProblems.LogDensityOrder{0}() +end +function LogDensityProblems.capabilities(::Type{<:VectorEvaluator}) + return LogDensityProblems.LogDensityOrder{0}() +end + +function LogDensityProblems.logdensity_and_gradient(p::Prepared{<:Any,<:VectorEvaluator}, x) + val, grad = AbstractPPL.value_and_gradient!!(p, x) + # `value_and_gradient!!` may alias internal storage; LDP requires a stable result. + return (val, copy(grad)) +end + +end # module diff --git a/src/AbstractPPL.jl b/src/AbstractPPL.jl index e08ee8b4..48a2cb44 100644 --- a/src/AbstractPPL.jl +++ b/src/AbstractPPL.jl @@ -10,6 +10,12 @@ export AbstractModelTrace include("abstractmodeltrace.jl") include("abstractprobprog.jl") include("evaluate.jl") +include("evaluators/Evaluators.jl") +using .Evaluators: prepare, value_and_gradient!!, value_and_jacobian!! +@static if VERSION >= v"1.11.0" + eval(Meta.parse("public prepare, value_and_gradient!!, value_and_jacobian!!")) +end + include("varname/optic.jl") include("varname/varname.jl") include("varname/subsumes.jl") diff --git a/src/evaluate.jl b/src/evaluate.jl index 57062062..19338565 100644 --- a/src/evaluate.jl +++ b/src/evaluate.jl @@ -5,7 +5,7 @@ Common base type for evaluation contexts. """ abstract type AbstractContext end -""" +""" evaluate!! General API for model operations, e.g. prior evaluation, log density, log joint etc. diff --git a/src/evaluators/Evaluators.jl b/src/evaluators/Evaluators.jl new file mode 100644 index 00000000..ff2d5407 --- /dev/null +++ b/src/evaluators/Evaluators.jl @@ -0,0 +1,263 @@ +module Evaluators + +using ADTypes: AbstractADType +import ..evaluate!! + +include("utils.jl") + +""" + Prepared{AD<:AbstractADType,E,C}(adtype, evaluator, cache) + Prepared(adtype, evaluator) # cache defaults to `nothing` + +AD-prepared evaluator parameterised by backend type `AD`. + +- `adtype` — the backend, used for dispatch. +- `evaluator` — the user-facing callable (typically a `VectorEvaluator` or + `NamedTupleEvaluator`); forwarded on `p(x)`. +- `cache` — backend-specific pre-allocated state (ForwardDiff configs, Mooncake + caches, DifferentiationInterface preps, etc.). `Nothing` when the backend requires + no cached state. + +Extension packages implement `value_and_gradient!!` (and optionally +`value_and_jacobian!!`) by specialising on the `cache` type: + +```julia +function AbstractPPL.value_and_gradient!!( + p::Prepared{<:AbstractADType, <:VectorEvaluator, <:MyCache}, x::AbstractVector +) + ... +end +``` +""" +struct Prepared{AD<:AbstractADType,E,C} + adtype::AD + evaluator::E + cache::C +end + +Prepared(adtype::AbstractADType, evaluator) = Prepared(adtype, evaluator, nothing) + +(p::Prepared)(x) = p.evaluator(x) + +""" + prepare(problem, values::NamedTuple; check_dims::Bool=true) + prepare(problem, x::AbstractVector{<:Real}; check_dims::Bool=true) + prepare(adtype, problem, x::AbstractVector{<:Real}; check_dims::Bool=true) + +Prepare a callable evaluator for `problem`. + +Use the two-argument form with a `NamedTuple` when the evaluator works with +named inputs, or with a vector when it works with vector inputs. The +three-argument form, contributed by AD-backend extensions, additionally +prepares gradient or jacobian machinery for vector inputs. + +`check_dims` (default `true`) controls whether the returned evaluator validates +the input shape on each call. Pass `check_dims=false` to skip the per-call +check, e.g. inside an AD backend's hot path where the input shape is already +guaranteed. +""" +function prepare end + +# Default: wrap the callable in the appropriate evaluator so per-call shape +# checks fire even without a backend-specific `prepare` method. Downstream +# packages (e.g. DynamicPPL) override these for their problem types. +function prepare(problem, values::NamedTuple; check_dims::Bool=true) + return NamedTupleEvaluator{check_dims}(problem, values) +end +function prepare(problem, x::AbstractVector{<:Real}; check_dims::Bool=true) + return VectorEvaluator{check_dims}(problem, length(x)) +end + +""" + value_and_gradient!!(prepared, x::AbstractVector{<:Real}) + +Return `(value, gradient)` for a scalar-valued evaluator, potentially reusing +internal cache buffers of `prepared`. The returned gradient may alias +`prepared`'s internal storage; copy if you need to retain it past the next call. +""" +function value_and_gradient!! end + +""" + value_and_jacobian!!(prepared, x::AbstractVector{<:Real}) + +Return `(value::AbstractVector, jacobian::AbstractMatrix)` for a vector-valued +evaluator, potentially reusing internal cache buffers. The returned arrays may +alias `prepared`'s internal storage; copy if needed. +The Jacobian has shape `(length(value), length(x))`. +""" +function value_and_jacobian!! end + +""" + VectorEvaluator{CheckInput}(f, dim) + VectorEvaluator(f, dim) # equivalent to `VectorEvaluator{true}(f, dim)` + +Evaluator shape for scalar functions of a vector input. Part of the extension +author API; end users interact with the wrapping `Prepared` instead. + +`CheckInput` controls whether each call validates the input length. The default +(`true`) is the safe shape exposed via `prepared(x)`. Pass `CheckInput=false` +(via `check_dims=false` in `prepare`) for the callable handed to AD libraries, +where input shape is already guaranteed and the runtime check would persist in +the dual/shadow hot path. + +A bare `VectorEvaluator` is *not* differentiable; gradient capability is the +contract of the wrapping `Prepared` returned by `prepare(adtype, ...)`. +""" +struct VectorEvaluator{CheckInput,F} + f::F + dim::Int + function VectorEvaluator{CheckInput}(f::F, dim::Int) where {CheckInput,F} + CheckInput isa Bool || throw(ArgumentError("`CheckInput` must be a Bool.")) + dim >= 0 || throw(ArgumentError("`dim` must be non-negative, got $dim.")) + return new{CheckInput,F}(f, dim) + end +end + +VectorEvaluator(f, dim::Int) = VectorEvaluator{true}(f, dim) + +""" + NamedTupleEvaluator{CheckInput}(f, inputspec) + NamedTupleEvaluator(f, inputspec) # equivalent to `NamedTupleEvaluator{true}(f, inputspec)` + +Evaluator shape for functions of a `NamedTuple` input with a stable prototype. +Part of the extension author API; end users interact with the wrapping `Prepared`. + +The `inputspec` prototype's leaves must be one of: + +- `Real` or `Complex` (scalar) +- `AbstractArray` whose elements are themselves supported leaves +- `Tuple` or `NamedTuple` recursively containing supported leaves + +This matches the structural model used by [`flatten_to!!`](@ref) / +[`unflatten_to!!`](@ref). Other leaf types (e.g. `String`, `Symbol`, custom +structs) trigger an `ArgumentError` from the per-call shape check. + +`CheckInput` controls whether each call validates that the input `NamedTuple` +matches the prototype's `typeof` and per-leaf array `size`. The default (`true`) +is the safe shape exposed via `prepared(x)`. Pass `CheckInput=false` (via +`check_dims=false` in `prepare`) for the callable handed to AD libraries: the +prototype's `typeof` is captured at preparation time using the original element +types, so a `CheckInput=true` evaluator will reject inputs whose leaves are +dual/shadow numbers (or any other widened element type) even when the structure +is otherwise correct. +""" +struct NamedTupleEvaluator{CheckInput,F,P<:NamedTuple} + f::F + inputspec::P + function NamedTupleEvaluator{CheckInput}( + f::F, inputspec::P + ) where {CheckInput,F,P<:NamedTuple} + CheckInput isa Bool || throw(ArgumentError("`CheckInput` must be a Bool.")) + return new{CheckInput,F,P}(f, inputspec) + end +end + +NamedTupleEvaluator(f, inputspec::NamedTuple) = NamedTupleEvaluator{true}(f, inputspec) + +# Reject integer vectors with a clear error rather than letting them flow into +# AD backends (which usually fail confusingly). `T <: Integer` resolves at +# compile time, so the AD hot path (Float/dual `T`) elides the branch entirely. +function _reject_integer_input(x) + throw( + ArgumentError( + "VectorEvaluator requires a vector of floating-point values, but received an `$(typeof(x))`. Convert to a floating-point vector (e.g. `Float64.(x)`) before calling.", + ), + ) +end + +function (e::VectorEvaluator{true})(x::AbstractVector{T}) where {T} + T <: Integer && _reject_integer_input(x) + length(x) == e.dim || throw( + DimensionMismatch( + "Expected a vector of length $(e.dim), but got length $(length(x))." + ), + ) + return e.f(x) +end + +function (e::VectorEvaluator{false})(x::AbstractVector{T}) where {T} + T <: Integer && _reject_integer_input(x) + return e.f(x) +end + +function (e::NamedTupleEvaluator{true})(values::NamedTuple) + _assert_namedtuple_shape(e, values) + return e.f(values) +end +(e::NamedTupleEvaluator{false})(values::NamedTuple) = e.f(values) + +""" + _assert_namedtuple_shape(e::NamedTupleEvaluator, values) + +Throw `ArgumentError` unless `values` has the same type as the prototype captured +during preparation, including matching `size` for any nested `AbstractArray` +leaves. Also throws if the prototype contains a leaf type outside the supported +set (`Real`, `Complex`, `AbstractArray`, `Tuple`, `NamedTuple`). No-op when `e` +was constructed with `CheckInput=false`. +""" +function _assert_namedtuple_shape(e::NamedTupleEvaluator{true}, values) + typeof(values) === typeof(e.inputspec) || throw( + ArgumentError( + "Expected the same NamedTuple structure that was used to prepare this evaluator.", + ), + ) + _shapes_match(values, e.inputspec) || throw( + ArgumentError( + "Nested array shape differs from the prototype captured during preparation." + ), + ) + return nothing +end +_assert_namedtuple_shape(::NamedTupleEvaluator{false}, _) = nothing + +# Complements the `typeof` check above: same-typed arrays can differ in `size`. +# Arrays with non-`Real`/`Complex` eltype are walked element-wise to catch +# inner mismatches. Unknown leaves throw, mirroring the supported-leaves +# contract of the flatten/unflatten utilities in `utils.jl`. +# +# `Tuple` recursion uses `first`/`Base.tail` rather than a `zip` loop so each +# leaf call sees concrete element types — same idiom as `_unflatten`. +_shapes_match(::Union{Real,Complex}, ::Union{Real,Complex}) = true +function _shapes_match(a::AbstractArray, b::AbstractArray) + size(a) == size(b) || return false + eltype(a) <: Union{Real,Complex} && return true + for (ai, bi) in zip(a, b) + _shapes_match(ai, bi) || return false + end + return true +end +_shapes_match(::Tuple{}, ::Tuple{}) = true +function _shapes_match(a::Tuple, b::Tuple) + _shapes_match(first(a), first(b)) || return false + return _shapes_match(Base.tail(a), Base.tail(b)) +end +_shapes_match(a::NamedTuple, b::NamedTuple) = _shapes_match(values(a), values(b)) +function _shapes_match(a, _) + throw( + ArgumentError( + "Cannot validate shape for prototype leaf of type `$(typeof(a))`. Supported leaves are `Real`, `Complex`, `AbstractArray`, `Tuple`, and `NamedTuple`.", + ), + ) +end + +# Make prepared evaluators usable through the same `evaluate!!` API as models. +evaluate!!(e::Union{Prepared,VectorEvaluator,NamedTupleEvaluator}, x) = e(x) + +function __init__() + Base.Experimental.register_error_hint(MethodError) do io, exc, args, kwargs + # `args` are argument types, not values (see `Base.Experimental.show_error_hints`). + # Only fire when no extension has registered any AD-aware `prepare` method yet — + # once a backend is loaded, the candidate list in the `MethodError` is more + # informative than a generic "load an extension" hint. + exc.f === prepare || return nothing + length(args) >= 1 && args[1] <: AbstractADType || return nothing + # `nargs` counts `self`, so `>= 4` matches the AD-aware 3-positional form. + any(m -> m.nargs >= 4, methods(prepare)) && return nothing + print( + io, + "\nCalling `prepare` with an AD backend requires loading the corresponding extension (e.g., `using DifferentiationInterface`).", + ) + end +end + +end # module diff --git a/src/evaluators/utils.jl b/src/evaluators/utils.jl new file mode 100644 index 00000000..5e97d16c --- /dev/null +++ b/src/evaluators/utils.jl @@ -0,0 +1,163 @@ +# Vectorisation utilities + +# Opt-in: only `Array` round-trips cleanly through `similar` + `copyto!` +# preserving `typeof`. `SubArray` is excluded because `similar(::SubArray)` +# returns a plain `Array`, silently breaking the typeof round-trip contract +# advertised by `unflatten_to!!`. Structured/lazy wrappers and `OffsetArray` +# fall through to the catch-all — callers must `collect` first. +# +# TODO: extend with proper support for structured arrays (independent-entry +# packing) and factorisation types (Cholesky in particular is needed for PPL +# covariance parameters). + +flat_length(x::Union{Real,Complex}) = 1 +flat_length(x::Array{<:Union{Real,Complex}}) = length(x) +flat_length(::Tuple{}) = 0 +flat_length(x::Tuple) = sum(flat_length, x) +flat_length(::NamedTuple{(),Tuple{}}) = 0 +flat_length(x::NamedTuple) = sum(flat_length, values(x)) +flat_length(x) = throw(ArgumentError("This value cannot be flattened into a vector.")) + +flat_eltype(x::Union{Real,Complex}) = typeof(x) +flat_eltype(x::Array{T}) where {T<:Union{Real,Complex}} = T +flat_eltype(::Tuple{}) = Float64 +flat_eltype(x::Tuple) = mapreduce(flat_eltype, promote_type, x) +flat_eltype(::NamedTuple{(),Tuple{}}) = Float64 +flat_eltype(x::NamedTuple) = mapreduce(flat_eltype, promote_type, values(x)) +flat_eltype(x) = throw(ArgumentError("This value cannot be flattened into a vector.")) + +""" + flatten_to!!(buf, x) + +Flatten `x` into the vector-like buffer `buf`. + +Supported `x` values are: +- `Real` +- `Complex` +- `Array{<:Union{Real,Complex}}` (other `AbstractArray` subtypes, including + views, must be `collect`ed first) +- `Tuple` recursively containing supported values +- `NamedTuple` recursively containing supported values + +Pass `nothing` as `buf` to allocate a new vector. +""" +function flatten_to!!(::Nothing, x) + buf = Vector{flat_eltype(x)}(undef, flat_length(x)) + _flatten_to!(buf, x, 1) + return buf +end + +function flatten_to!!(buf::AbstractVector, x) + Base.require_one_based_indexing(buf) + n = flat_length(x) + length(buf) == n || throw( + DimensionMismatch("Expected a vector of length $n, but got length $(length(buf))."), + ) + _flatten_to!(buf, x, 1) + return buf +end + +function _flatten_to!(buf::AbstractVector, x::Union{Real,Complex}, offset::Int) + buf[offset] = x + return offset + 1 +end + +function _flatten_to!(buf::AbstractVector, x::Array{<:Union{Real,Complex}}, offset::Int) + Base.require_one_based_indexing(x) + n = length(x) + copyto!(buf, offset, x, 1, n) + return offset + n +end + +function _flatten_to!(buf::AbstractVector, x::Tuple, offset::Int) + for value in x + offset = _flatten_to!(buf, value, offset) + end + return offset +end + +function _flatten_to!(buf::AbstractVector, x::NamedTuple, offset::Int) + for value in values(x) + offset = _flatten_to!(buf, value, offset) + end + return offset +end + +function _flatten_to!(buf::AbstractVector, x, ::Int) + throw(ArgumentError("This value cannot be flattened into a vector.")) +end + +function _unflatten(x::Union{Real,Complex}, buf::AbstractVector, offset::Int) + return convert(typeof(x), buf[offset]), offset + 1 +end + +function _unflatten(x::Array{<:Union{Real,Complex}}, buf::AbstractVector, offset::Int) + Base.require_one_based_indexing(x) + n = length(x) + value = similar(x) + copyto!(value, 1, buf, offset, n) + return value, offset + n +end + +_unflatten(::Tuple{}, buf::AbstractVector, offset::Int) = (), offset + +function _unflatten(x::Tuple, buf::AbstractVector, offset::Int) + first_value, offset = _unflatten(first(x), buf, offset) + rest_value, offset = _unflatten(Base.tail(x), buf, offset) + return (first_value, rest_value...), offset +end + +# Generated to keep the result `NamedTuple` type inferable: a recursive `merge` +# over `Base.tail(Names)` erases parameters and breaks `@inferred` callers. +@generated function _unflatten( + x::NamedTuple{Names}, buf::AbstractVector, offset::Int +) where {Names} + if isempty(Names) + return :((NamedTuple(), offset)) + end + block = Expr(:block, :(off = offset)) + val_syms = Symbol[] + for name in Names + v = gensym(name) + push!(val_syms, v) + push!(block.args, :(($v, off) = _unflatten(x[$(QuoteNode(name))], buf, off))) + end + push!(block.args, :(return (NamedTuple{$Names}(($(val_syms...),)), off))) + return block +end + +""" + unflatten_to!!(x, buf; check_eltype::Bool=false) + +Reconstruct a value from the vector-like buffer `buf` using `x` as the structural template. + +Supported `x` values are: +- `Real` +- `Complex` +- `Array{<:Union{Real,Complex}}` (other `AbstractArray` subtypes, including + views, must be `collect`ed first) +- `Tuple` recursively containing supported values +- `NamedTuple` recursively containing supported values + +Pass `check_eltype=true` to emit a warning when `eltype(buf)` differs from +`flat_eltype(x)` (off by default to keep hot paths quiet). + +Leaves are rebuilt using `x`'s types, so `typeof(unflatten_to!!(x, buf)) == typeof(x)` +even when `buf`'s element type is widened (e.g. real `x` flattened into a `ComplexF64` +buffer). Always allocates: each array leaf goes through `similar`. +""" +function unflatten_to!!(x, buf::AbstractVector; check_eltype::Bool=false) + Base.require_one_based_indexing(buf) + n = flat_length(x) + length(buf) == n || throw( + DimensionMismatch("Expected a vector of length $n, but got length $(length(buf))."), + ) + if check_eltype + expected = flat_eltype(x) + eltype(buf) === expected || @warn( + "Buffer eltype `$(eltype(buf))` differs from `flat_eltype(x) = $expected`; reconstructing using the leaf types from `x`. An `InexactError` will be thrown if any value in `buf` cannot be converted back to the corresponding leaf type." + ) + end + value, _ = _unflatten(x, buf, 1) + return value +end diff --git a/test/Project.toml b/test/Project.toml index a43fb829..fb93b3bd 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,4 +1,6 @@ [deps] +ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" +AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf" Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697" Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" DimensionalData = "0703355e-b756-11e9-17c0-8b28908087d0" @@ -13,6 +15,7 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [compat] +ADTypes = "1" Accessors = "0.1" Aqua = "0.8" DimensionalData = "0.29, 0.30" diff --git a/test/evaluators/Evaluators.jl b/test/evaluators/Evaluators.jl new file mode 100644 index 00000000..c8454438 --- /dev/null +++ b/test/evaluators/Evaluators.jl @@ -0,0 +1,149 @@ +using AbstractPPL +using AbstractPPL: prepare, value_and_gradient!!, evaluate!! +using AbstractPPL.Evaluators: Prepared, VectorEvaluator, NamedTupleEvaluator +using ADTypes: ADTypes +using Test + +struct DummyModel end + +struct DummyPrepared + prototype_keys::Tuple +end + +function AbstractPPL.prepare(model::DummyModel, values::NamedTuple) + return DummyPrepared(keys(values)) +end + +function (p::DummyPrepared)(values::NamedTuple) + keys(values) == p.prototype_keys || + error("expected fields $(p.prototype_keys), got $(keys(values))") + return sum(x -> x isa AbstractArray ? sum(x) : x, values) +end + +struct DummyADType <: ADTypes.AbstractADType end + +function AbstractPPL.prepare( + adtype::DummyADType, model::DummyModel, x::AbstractVector{<:Real}; check_dims::Bool=true +) + f = x -> sum(x) + return Prepared(adtype, VectorEvaluator{check_dims}(f, length(x))) +end + +function AbstractPPL.value_and_gradient!!( + p::Prepared{DummyADType}, x::AbstractVector{<:Real} +) + return (sum(x), ones(length(x))) +end + +@testset "Evaluators interface" begin + @testset "explicit evaluator shapes" begin + ve = AbstractPPL.Evaluators.VectorEvaluator(sum, 3) + @test ve([1.0, 2.0, 3.0]) == 6.0 + @test_throws DimensionMismatch ve([1.0, 2.0]) + @test_throws r"floating-point" ve([1, 2, 3]) + + ne = AbstractPPL.Evaluators.NamedTupleEvaluator( + x -> x.a + sum(x.b), (a=0.0, b=zeros(2)) + ) + @test ne((a=1.0, b=[2.0, 3.0])) == 6.0 + @test ne.inputspec == (a=0.0, b=zeros(2)) + @test_throws MethodError ne([1.0, 2.0, 3.0]) + + # `CheckInput=false` skips the per-call shape checks. + ve_unchecked = AbstractPPL.Evaluators.VectorEvaluator{false}(sum, 3) + @test ve_unchecked([1.0, 2.0]) == 3.0 + + ne_unchecked = AbstractPPL.Evaluators.NamedTupleEvaluator{false}( + x -> 0.0, (a=0.0, b=zeros(2)) + ) + @test ne_unchecked((totally=:wrong,)) == 0.0 + @test_throws r"same NamedTuple structure" ne((totally=:wrong,)) + + # Nested array shape: same `typeof` (Vector{Float64}), different size. + @test_throws r"Nested array" ne((a=1.0, b=[2.0])) + + # Array-of-arrays: same `typeof` and outer size, mismatched inner size. + ne_nested = AbstractPPL.Evaluators.NamedTupleEvaluator( + x -> sum(sum, x.b), (b=[zeros(2), zeros(2)],) + ) + @test ne_nested((b=[[1.0, 2.0], [3.0, 4.0]],)) == 10.0 + @test_throws r"Nested array" ne_nested((b=[[1.0], [2.0]],)) + + # Unsupported leaf types are rejected rather than silently passing. + ne_string = AbstractPPL.Evaluators.NamedTupleEvaluator(x -> length(x.s), (s="abc",)) + @test_throws r"Supported leaves" ne_string((s="abcde",)) + end + + @testset "prepare (structural)" begin + model = DummyModel() + values = (x=0.0, y=[1.0, 2.0]) + prepared = prepare(model, values) + @test prepared isa DummyPrepared + @test prepared.prototype_keys == (:x, :y) + + lp = prepared((x=0.5, y=[1.5, 2.5])) + @test lp ≈ 0.5 + 1.5 + 2.5 + + @test_throws ErrorException prepared((a=1.0, b=2.0)) + + # Generic fallback wraps a plain callable in the appropriate evaluator + # so per-call shape checks fire even without a backend-specific override. + pv = prepare(sum, zeros(3)) + @test pv isa VectorEvaluator{true} + @test pv([1.0, 2.0, 3.0]) == 6.0 + @test_throws DimensionMismatch pv([1.0, 2.0]) + + ntfun = v -> v.a + sum(v.b) + pn = prepare(ntfun, (a=0.0, b=zeros(2))) + @test pn isa NamedTupleEvaluator{true} + @test pn((a=1.0, b=[2.0, 3.0])) == 6.0 + + # check_dims=false propagates to the wrapper. + pv_unchecked = prepare(sum, zeros(3); check_dims=false) + @test pv_unchecked isa VectorEvaluator{false} + @test pv_unchecked([1.0, 2.0]) == 3.0 # wrong length, no error + end + + @testset "prepare (AD-aware)" begin + model = DummyModel() + x0 = zeros(3) + adtype = DummyADType() + prepared = prepare(adtype, model, x0) + @test prepared isa Prepared{DummyADType} + + x = [0.5, 1.5, 2.5] + @test prepared(x) ≈ 0.5 + 1.5 + 2.5 + + val, grad = value_and_gradient!!(prepared, x) + @test val ≈ 0.5 + 1.5 + 2.5 + @test grad ≈ [1.0, 1.0, 1.0] + + # check_dims=false skips the per-call dimension check. + prepared_unchecked = prepare(adtype, model, x0; check_dims=false) + @test prepared_unchecked([1.0, 2.0]) ≈ 3.0 # wrong length, no error + end + + @testset "missing AD package extensions" begin + model = DummyModel() + x0 = zeros(3) + + @test_throws MethodError AbstractPPL.Evaluators.prepare( + ADTypes.AutoEnzyme(), model, x0 + ) + end + + @testset "evaluate!!" begin + ve = AbstractPPL.Evaluators.VectorEvaluator(sum, 3) + @test evaluate!!(ve, [1.0, 2.0, 3.0]) == 6.0 + @test_throws DimensionMismatch evaluate!!(ve, [1.0, 2.0]) + + ne = AbstractPPL.Evaluators.NamedTupleEvaluator( + x -> x.a + sum(x.b), (a=0.0, b=zeros(2)) + ) + @test evaluate!!(ne, (a=1.0, b=[2.0, 3.0])) == 6.0 + + adtype = DummyADType() + prepared = prepare(adtype, DummyModel(), zeros(3)) + @test evaluate!!(prepared, [0.5, 1.5, 2.5]) ≈ 4.5 + end +end diff --git a/test/evaluators/utils.jl b/test/evaluators/utils.jl new file mode 100644 index 00000000..fb7f8571 --- /dev/null +++ b/test/evaluators/utils.jl @@ -0,0 +1,115 @@ +using AbstractPPL +using AbstractPPL.Evaluators: flatten_to!!, unflatten_to!! +using LinearAlgebra: Symmetric +using OffsetArrays: OffsetArray +using Test + +@testset "vectorisation utilities" begin + @testset "scalar round-trip" begin + x = 1.5 + v = flatten_to!!(nothing, x) + @test v == [1.5] + @test unflatten_to!!(x, v) == 1.5 + + z = 1.0 + 2.0im + vz = flatten_to!!(nothing, z) + @test vz == ComplexF64[1.0 + 2.0im] + @test unflatten_to!!(z, vz) == z + end + + @testset "array round-trip" begin + x = [1.0 2.0; 3.0 4.0] + v = flatten_to!!(nothing, x) + @test v == [1.0, 3.0, 2.0, 4.0] + @test unflatten_to!!(x, v) == x + + z = ComplexF64[1.0 + 1.0im, 2.0 + 0.0im] + vz = flatten_to!!(nothing, z) + @test vz == z + @test unflatten_to!!(z, vz) == z + end + + @testset "tuple round-trip" begin + x = (1.0, [2.0, 3.0], (4.0 + 1.0im,)) + v = flatten_to!!(nothing, x) + @test v == ComplexF64[1.0, 2.0, 3.0, 4.0 + 1.0im] + @test unflatten_to!!(x, v) == x + end + + @testset "named tuple round-trip" begin + x = (a=1.0, b=([2.0, 3.0], (c=4.0 + 1.0im,))) + v = flatten_to!!(nothing, x) + @test v == ComplexF64[1.0, 2.0, 3.0, 4.0 + 1.0im] + @test unflatten_to!!(x, v) == x + end + + @testset "heterogeneous container preserves leaf types" begin + # The flat buffer widens to ComplexF64, but `unflatten_to!!` rebuilds + # leaves using `x`'s types, so the round-trip preserves `typeof(x)`. + x = (1.0, [2.0, 3.0], (4.0 + 1.0im,)) + x2 = unflatten_to!!(x, flatten_to!!(nothing, x)) + @test x2 == x + @test typeof(x2) == typeof(x) + end + + @testset "check_eltype opt-in warning" begin + x = (a=1.0, b=[2.0, 3.0]) + # buf eltype is ComplexF64, but `flat_eltype(x) == Float64` — a + # mismatch that should warn only with `check_eltype=true`. Imag parts + # are zero, so the convert succeeds. + buf = ComplexF64[1.0, 2.0, 3.0] + x2 = @test_logs unflatten_to!!(x, buf) # default: silent + @test x2 == x + @test typeof(x2) == typeof(x) + x3 = @test_logs (:warn, r"differs from") unflatten_to!!(x, buf; check_eltype=true) + @test x3 == x + @test typeof(x3) == typeof(x) + end + + @testset "non-one-based buf rejected" begin + # Non-one-based `buf` fails `require_one_based_indexing` even when `x` + # is a plain `Array`. (Non-one-based `x` is covered by the catch-all + # in the next testset.) + @test_throws ArgumentError flatten_to!!(OffsetArray(zeros(3), 0:2), [1.0, 2.0, 3.0]) + @test_throws ArgumentError unflatten_to!!( + [1.0, 2.0, 3.0], OffsetArray(zeros(3), 0:2) + ) + end + + @testset "non-Array AbstractArrays rejected" begin + for x in ( + Symmetric([1.0 2.0; 2.0 3.0]), + @view([1.0, 2.0, 3.0, 4.0][1:4]), + OffsetArray([1.0, 2.0, 3.0], 0:2), + ) + @test_throws r"cannot be flattened" flatten_to!!(nothing, x) + @test_throws r"cannot be flattened" unflatten_to!!(x, [1.0, 2.0, 3.0, 4.0]) + end + end + + @testset "buffer length mismatch" begin + @test_throws r"Expected a vector of length 4" flatten_to!!( + Vector{Float64}(undef, 3), zeros(2, 2) + ) + end + + @testset "vector length mismatch" begin + x = (a=1.0, b=[2.0, 3.0]) + @test_throws r"Expected a vector of length 3" unflatten_to!!(x, [1.0, 2.0]) + end + + @testset "edge cases" begin + empty = NamedTuple() + @test flatten_to!!(nothing, empty) == Float64[] + @test unflatten_to!!(empty, Float64[]) == empty + end + + @testset "unflatten_to!! type stability" begin + @inferred unflatten_to!!((a=1.0, b=2.0, c=3.0), zeros(3)) + @inferred unflatten_to!!((a=1.0, b=[2.0, 3.0], c=3.0), zeros(4)) + @inferred unflatten_to!!((a=(p=1.0, q=2.0), b=3.0), zeros(3)) + @inferred unflatten_to!!((a=1.0, b=(2.0, 3.0)), zeros(3)) + @inferred unflatten_to!!(NamedTuple(), Float64[]) + @inferred unflatten_to!!((1.0, [2.0, 3.0], 3.0), zeros(4)) + end +end diff --git a/test/ext/logdensityproblems/Project.toml b/test/ext/logdensityproblems/Project.toml new file mode 100644 index 00000000..a6d089ab --- /dev/null +++ b/test/ext/logdensityproblems/Project.toml @@ -0,0 +1,11 @@ +[deps] +AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf" +ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" +LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c" +Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" + +[compat] +ADTypes = "1" +LogDensityProblems = "2" +julia = "1.10" diff --git a/test/ext/logdensityproblems/main.jl b/test/ext/logdensityproblems/main.jl new file mode 100644 index 00000000..c0176035 --- /dev/null +++ b/test/ext/logdensityproblems/main.jl @@ -0,0 +1,70 @@ +using Pkg +Pkg.activate(@__DIR__) +Pkg.develop(; path=joinpath(@__DIR__, "..", "..", "..")) +Pkg.instantiate() + +using AbstractPPL +using AbstractPPL.Evaluators: Prepared, VectorEvaluator, NamedTupleEvaluator +using ADTypes: AbstractADType, AutoForwardDiff +using LogDensityProblems: LogDensityProblems +using Test + +# A NamedTupleEvaluator does not satisfy LDP's vector-input contract, so the +# extension does not define LDP methods for it. + +struct TestADType <: AbstractADType end + +function AbstractPPL.value_and_gradient!!( + p::Prepared{TestADType}, x::AbstractVector{<:Real} +) + return (p(x), ones(length(x))) +end + +# Backend extensions opt into gradient capability by overloading `capabilities` +# (typically on their cache type, e.g. `<:Prepared{<:Any,<:VectorEvaluator,<:MyCache}`). +# Here we dispatch on the AD type for simplicity. +function LogDensityProblems.capabilities(::Type{<:Prepared{TestADType,<:VectorEvaluator}}) + return LogDensityProblems.LogDensityOrder{1}() +end + +@testset "AbstractPPLLogDensityProblemsExt" begin + @testset "VectorEvaluator" begin + ve = VectorEvaluator(sum, 3) + @test LogDensityProblems.dimension(ve) == 3 + @test LogDensityProblems.logdensity(ve, [1.0, 2.0, 3.0]) == 6.0 + # A bare VectorEvaluator never advertises gradient capability; + # only the wrapping `Prepared` does. + @test LogDensityProblems.capabilities(ve) == LogDensityProblems.LogDensityOrder{0}() + end + + @testset "Prepared capabilities" begin + # Without a backend overload the fallback advertises order 0 only. + p_no_overload = Prepared(AutoForwardDiff(), VectorEvaluator(sum, 3)) + @test LogDensityProblems.capabilities(p_no_overload) == + LogDensityProblems.LogDensityOrder{0}() + + # A backend that overloads capabilities advertises order 1. + p_overloaded = Prepared(TestADType(), VectorEvaluator(sum, 3)) + @test LogDensityProblems.capabilities(p_overloaded) == + LogDensityProblems.LogDensityOrder{1}() + @test LogDensityProblems.capabilities(typeof(p_overloaded)) == + LogDensityProblems.LogDensityOrder{1}() + + # NamedTupleEvaluator-backed Prepared has no LDP methods defined; the + # extension only integrates vector-input evaluators. + p_nt = Prepared( + AutoForwardDiff(), NamedTupleEvaluator(x -> x.a + sum(x.b), (a=0.0, b=zeros(2))) + ) + @test_throws MethodError LogDensityProblems.dimension(p_nt) + @test LogDensityProblems.capabilities(p_nt) === nothing + end + + @testset "logdensity_and_gradient" begin + f = x -> -0.5 * sum(abs2, x) + p = Prepared(TestADType(), VectorEvaluator(f, 3)) + x = [1.0, 2.0, 3.0] + val, grad = LogDensityProblems.logdensity_and_gradient(p, x) + @test val ≈ f(x) + @test grad ≈ ones(3) + end +end diff --git a/test/run_extras.jl b/test/run_extras.jl new file mode 100644 index 00000000..0e71bbcd --- /dev/null +++ b/test/run_extras.jl @@ -0,0 +1,12 @@ +# Run a named extension test in its own isolated Julia environment. +# +# Usage (from the repo root): +# LABEL=ext/logdensityproblems julia test/run_extras.jl + +const VALID_LABELS = ("ext/logdensityproblems",) + +label = get(ENV, "LABEL", nothing) +label in VALID_LABELS || + error("Set LABEL to one of: $(join(VALID_LABELS, ", ")) (got `$label`).") + +include(joinpath(@__DIR__, label, "main.jl")) diff --git a/test/runtests.jl b/test/runtests.jl index fc8d3980..ad97e108 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -8,6 +8,8 @@ const GROUP = get(ENV, "GROUP", "All") if GROUP == "All" || GROUP == "Tests" include("Aqua.jl") include("abstractprobprog.jl") + include("evaluators/Evaluators.jl") + include("evaluators/utils.jl") include("varname/optic.jl") include("varname/varname.jl") include("varname/subsumes.jl")