Move VarNamedTuple from DynamicPPL into AbstractPPL

HISTORY.md

@@ -1,3 +1,22 @@
+## 0.16.0
+
+Moved `VarNamedTuple` and its structured-array support from DynamicPPL into AbstractPPL so
+probabilistic programming frontends can share a lossless, `VarName`-keyed parameter
+container. The exported API includes `VarNamedTuple`, `@vnt`, `subset`, and the existing
+mapping, densification, and skeleton helpers. Template-based insertion remains public as
+`templated_setindex!!`. ComponentArrays integration is provided by a package extension, and
+the `Distribution`-aware `hasvalue`/`getvalue` methods live in the existing Distributions
+package extension.
+
+This release is purely additive over 0.15 — nothing was removed or changed in the existing
+public API, so downstream packages can widen their compat bound without code changes.
+However, the ten newly exported names (`VarNamedTuple`, `@vnt`, `map_pairs!!`,
+`map_values!!`, `apply!!`, `densify!!`, `skeleton`, `subset`, `NoTemplate`, `SkipTemplate`)
+are currently also exported by DynamicPPL for its own copy of the same functionality. Code
+that does `using DynamicPPL, AbstractPPL` unqualified will therefore see ambiguous bindings
+for these names until DynamicPPL migrates to re-export AbstractPPL's implementation;
+qualify uses in the meantime.
+
 ## 0.15.4
 
 `value_and_gradient!!`, `value_and_jacobian!!`, and `value_gradient_and_hessian!!` now accept a `context=` keyword that overrides, for a single call, the `context` frozen at `prepare` — without re-preparing (#167). The default (`context=nothing`) leaves the hot path unchanged; a `Tuple` override must match the frozen context's element types and shapes — a type mismatch (or a non-`Tuple` override) throws an `ArgumentError` — and reuses the type-keyed cache. Backends that bake context into their prepared state throw an `ArgumentError` for a non-empty-input override (re-`prepare` instead): compiled-tape ReverseDiff (`AutoReverseDiff(; compile=true)`) on all three entry points, and Mooncake's Hessian; Mooncake Jacobian preps are context-free by construction, so only an empty override validates there. Empty input runs no derivative machinery, so no backend rejects an override there (it is still validated).

Project.toml

@@ -3,7 +3,7 @@ uuid = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf"
 keywords = ["probabilistic programming"]
 license = "MIT"
 desc = "Common interfaces for probabilistic programming"
-version = "0.15.4"
+version = "0.16.0"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -19,6 +19,7 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 
 [weakdeps]
+ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5"
 DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
@@ -27,6 +28,7 @@ Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [extensions]
+AbstractPPLComponentArraysExt = ["ComponentArrays"]
 AbstractPPLDifferentiationInterfaceExt = ["DifferentiationInterface"]
 AbstractPPLDistributionsExt = ["Distributions", "LinearAlgebra"]
 AbstractPPLForwardDiffExt = ["ForwardDiff", "DiffResults"]
@@ -37,7 +39,8 @@ AbstractPPLTestExt = ["Test"]
 ADTypes = "1"
 AbstractMCMC = "2, 3, 4, 5"
 Accessors = "0.1"
-BangBang = "0.4"
+BangBang = "0.4.1"
+ComponentArrays = "0.15"
 DensityInterface = "0.4"
 DiffResults = "1"
 DifferentiationInterface = "0.6, 0.7"

docs/make.jl

@@ -10,7 +10,15 @@ DocMeta.setdocmeta!(AbstractPPL, :DocTestSetup, :(using AbstractPPL); recursive=
 makedocs(;
     sitename="AbstractPPL",
     modules=[AbstractPPL, Base.get_extension(AbstractPPL, :AbstractPPLDistributionsExt)],
-    pages=["index.md", "varname.md", "of.md", "pplapi.md", "evaluators.md", "interface.md"],
+    pages=[
+        "index.md",
+        "varname.md",
+        "varnamedtuple.md",
+        "of.md",
+        "pplapi.md",
+        "evaluators.md",
+        "interface.md",
+    ],
     checkdocs=:exports,
     doctest=false,
 )

docs/src/varnamedtuple.md

@@ -0,0 +1,34 @@
+# VarNamedTuple
+
+`VarNamedTuple` is a lossless, `VarName`-keyed value container. It preserves nested
+property and array structure while allowing probabilistic programming packages to exchange
+values without depending on one frontend's model representation.
+
+Use [`@vnt`](@ref) for literal construction and a `VarName` to access a value. Functional
+updates are available through BangBang's `setindex!!` interface.
+
+```@example varnamedtuple
+using AbstractPPL
+
+values = @vnt begin
+    location := 1.0
+    scale := 2.0
+end
+values[@varname(location)]
+```
+
+## API
+
+```@docs
+VarNamedTuple
+AbstractPPL.@vnt
+subset
+map_pairs!!
+map_values!!
+apply!!
+densify!!
+skeleton
+NoTemplate
+SkipTemplate
+AbstractPPL.templated_setindex!!
+```

ext/AbstractPPLComponentArraysExt.jl

@@ -0,0 +1,55 @@
+module AbstractPPLComponentArraysExt
+using AbstractPPL: AbstractPPL
+using AbstractPPL.VarNamedTuples:
+    PartialArray, AllowAll, SetPermissions, _setindex_optic!!, _getindex_optic, make_leaf
+using ComponentArrays: ComponentArrays, ComponentVector
+
+# Helper: convert a Property optic label S to an integer Index optic
+function _property_to_index(
+    template::ComponentVector, optic::AbstractPPL.Property{S}
+) where {S}
+    ax = ComponentArrays.getaxes(template)[1]
+    idx = first(ax[S].idx)
+    return AbstractPPL.Index((idx,), NamedTuple(), optic.child)
+end
+
+function AbstractPPL.VarNamedTuples.make_leaf(
+    value, optic::AbstractPPL.Property{S}, template::ComponentVector
+) where {S}
+    return if optic.child isa AbstractPPL.Iden
+        index_optic = _property_to_index(template, optic)
+        make_leaf(value, index_optic, template)
+    else
+        # This branch is needed to handle nested axes in ComponentArrays: the idea is that
+        # if x is e.g. ComponentArray(a=(b=1)) and we are trying to set `x.a.b`, then we
+        # first index into `x.a` to get the slice of the ComponentArray. The easiest way to
+        # handle this is to call the default method.
+        invoke(
+            make_leaf,
+            Tuple{Any,AbstractPPL.Property{S},AbstractArray},
+            value,
+            optic,
+            template,
+        )
+    end
+end
+
+function AbstractPPL.VarNamedTuples._setindex_optic!!(
+    pa::PartialArray{<:Any,<:Any,<:ComponentVector},
+    value,
+    optic::AbstractPPL.Property{S},
+    template,
+    permissions::SetPermissions=AllowAll(),
+) where {S}
+    index_optic = _property_to_index(pa.data, optic)
+    return _setindex_optic!!(pa, value, index_optic, template, permissions)
+end
+
+function AbstractPPL.VarNamedTuples._getindex_optic(
+    pa::PartialArray{<:Any,<:Any,<:ComponentVector}, optic::AbstractPPL.Property{S}, orig_vn
+) where {S}
+    index_optic = _property_to_index(pa.data, optic)
+    return _getindex_optic(pa, index_optic, orig_vn)
+end
+
+end

ext/AbstractPPLDistributionsExt.jl

@@ -50,6 +50,7 @@ This decision may be revisited in the future.
 module AbstractPPLDistributionsExt
 
 using AbstractPPL: AbstractPPL, VarName, Accessors
+using AbstractPPL.VarNamedTuples: VarNamedTuple, PartialArray, _getindex_optic
 using Distributions: Distributions
 using LinearAlgebra: Cholesky, LowerTriangular, UpperTriangular
 
@@ -70,7 +71,7 @@ function get_optics(dist::Distributions.LKJCholesky)
     is_up = dist.uplo == 'U'
     cartesian_indices = filter(CartesianIndices(size(dist))) do cartesian_index
         i, j = cartesian_index.I
-        is_up ? i <= j : i >= j
+        return is_up ? i <= j : i >= j
     end
     # there is an additional layer as we need to access `.L` or `.U` before we
     # can index into it
@@ -302,4 +303,61 @@ function AbstractPPL.getvalue(
     end
 end
 
+# A VarNamedTuple already retains its structure, so most distributions can use its direct
+# lookup. LKJCholesky is the exception: values may be stored under the `.L` or `.U` optics
+# and need the existing dictionary reconstruction path below.
+function AbstractPPL.hasvalue(
+    vals::VarNamedTuple, vn::VarName, ::Distributions.Distribution
+)
+    return haskey(vals, vn)
+end
+function AbstractPPL.getvalue(
+    vals::VarNamedTuple, vn::VarName, ::Distributions.Distribution
+)
+    return vals[vn]
+end
+
+function AbstractPPL.hasvalue(
+    vnt::VarNamedTuple, vn::VarName, dist::Distributions.LKJCholesky
+)
+    haskey(vnt, vn) || return false
+    val = _getindex_optic(vnt, vn)
+    if !(val isa VarNamedTuple || val isa PartialArray)
+        return true
+    end
+
+    et = val isa VarNamedTuple ? Any : eltype(val)
+    dval = Dict{VarName,et}()
+    for key in keys(val)
+        subvn = if val isa VarNamedTuple
+            AbstractPPL.prefix(key, vn)
+        else
+            AbstractPPL.append_optic(vn, key)
+        end
+        dval[subvn] = _getindex_optic(val, key, subvn)
+    end
+    return AbstractPPL.hasvalue(dval, vn, dist)
+end
+
+function AbstractPPL.getvalue(
+    vnt::VarNamedTuple, vn::VarName, dist::Distributions.LKJCholesky
+)
+    val = _getindex_optic(vnt, vn)
+    if !(val isa VarNamedTuple || val isa PartialArray)
+        return val
+    end
+
+    et = val isa VarNamedTuple ? Any : eltype(val)
+    dval = Dict{VarName,et}()
+    for key in keys(val)
+        subvn = if val isa VarNamedTuple
+            AbstractPPL.prefix(key, vn)
+        else
+            AbstractPPL.append_optic(vn, key)
+        end
+        dval[subvn] = _getindex_optic(val, key, subvn)
+    end
+    return AbstractPPL.getvalue(dval, vn, dist)
+end
+
 end

src/AbstractPPL.jl

@@ -101,6 +101,42 @@ export AbstractOptic,
 using Accessors: set
 export set
 
+"""
+    subset(value, varnames)
+
+Return the part of `value` selected by `varnames`.
+"""
+function subset end
+
+include("varnamedtuple.jl")
+using .VarNamedTuples:
+    VarNamedTuples,
+    VarNamedTuple,
+    map_pairs!!,
+    map_values!!,
+    apply!!,
+    densify!!,
+    templated_setindex!!,
+    NoTemplate,
+    SkipTemplate,
+    @vnt,
+    skeleton
+
+export VarNamedTuple,
+    @vnt,
+    map_pairs!!,
+    map_values!!,
+    apply!!,
+    densify!!,
+    skeleton,
+    subset,
+    NoTemplate,
+    SkipTemplate
+
+@static if VERSION >= v"1.11.0"
+    eval(Meta.parse("public VarNamedTuples, templated_setindex!!"))
+end
+
 include("of.jl")
 export of, @of
 @static if VERSION >= v"1.11.0"