VNT Part 2: Solving the issue of block elements (#1180)

* ArrayLikeBlock WIP

* ArrayLikeBlock WIP2

* Improve type stability of ArrayLikeBlock stuff

* Test more invariants

* Actually run VNT tests

* Implement show for ArrayLikeBlock

* Change keys on VNT to return an array

* Fix keys and some tests for PartialArray

* Improve type stability

* Fix keys for PartialArray

* More ArrayLikeBlock tests

* Add docstrings

* Remove redundant code, improve documentation

* Add Base.size(::RangeAndLinked)

* Fix issues with RangeAndLinked and VNT

* Write more design doc for ArrayLikeBlocks

Author: mhauru

docs/src/internals/varnamedtuple.md

@@ -144,6 +144,29 @@ You can also set the elements with `vnt = setindex!!(vnt, @varname(a[1]), 3.0)`,
 At this point you can not set any new values in that array that would be outside of its range, with something like `vnt = setindex!!(vnt, @varname(a[5]), 5.0)`.
 The philosophy here is that once a `Base.Array` has been attached to a `VarName`, that takes precedence, and a `PartialArray` is only used as a fallback when we are told to store a value for `@varname(a[i])` without having any previous knowledge about what `@varname(a)` is.
 
+## Non-Array blocks with `IndexLens`es
+
+The above is all that is needed for setting regular scalar values.
+However, in DynamicPPL we also have a particular need for something slightly odd:
+We sometimes need to do calls like `setindex!!(vnt, @varname(a[1:5]), val)` on a `val` that is _not_ an `AbstractArray`, or even iterable at all.
+Normally this would error: As a scalar value with size `()`, `val` is the wrong size to be set with `@varname(a[1:5])`, which clearly wants something with size `(5,)`.
+However, we want to allow this even if `val` is not an iterable, if it is some object for which `size` is well-defined, and `size(val) == (5,)`.
+In DynamicPPL this comes up when storing e.g. the priors of a model, where a random variable like `@varname(a[1:5])` may be associated with a prior that is a 5-dimensional distribution.
+
+Internally, a `PartialArray` is just a regular `Array` with a mask saying which elements have been set.
+Hence we can't store `val` directly in the same `PartialArray`:
+We need it to take up a sub-block of the array, in our example case a sub-block of length 5.
+To this end, internally, `PartialArray` uses a wrapper type called `ArrayLikeWrapper`, that stores `val` together with the indices that are being used to set it.
+The `PartialArray` has all its corresponding elements, in our example elements 1, 2, 3, 4, and, 5, point to the same wrapper object.
+
+While such blocks can be stored using a wrapper like this, some care must be taken in indexing into these blocks.
+For instance, after setting a block with `setindex!!(vnt, @varname(a[1:5]), val)`, we can't `getindex(vnt, @varname(a[1]))`, since we can't return "the first element of five in `val`", because `val` may not be indexable in any way.
+Similarly, if next we set `setindex!!(vnt, @varname(a[1]), some_other_value)`, that should invalidate/delete the elements `@varname(a[2:5])`, since the block only makes sense as a whole.
+Because of these reasons, setting and getting blocks of well-defined size like this is allowed with `VarNamedTuple`s, but _only by always using the full range_.
+For instance, if `setindex!!(vnt, @varname(a[1:5]), val)` has been set, then the only valid `getindex` key to access `val` is `@varname(a[1:5])`;
+Not `@varname(a[1:10])`, nor `@varname(a[3])`, nor for anything else that overlaps with `@varname(a[1:5])`.
+`haskey` likewise only returns true for `@varname(a[1:5])`, and `keys(vnt)` only has that as an element.
+
 ## Limitations
 
 This design has a several of benefits, for performance and generality, but it also has limitations:

ext/DynamicPPLMarginalLogDensitiesExt.jl

@@ -1,6 +1,6 @@
 module DynamicPPLMarginalLogDensitiesExt
 
-using DynamicPPL: DynamicPPL, LogDensityProblems, VarName
+using DynamicPPL: DynamicPPL, LogDensityProblems, VarName, RangeAndLinked
 using MarginalLogDensities: MarginalLogDensities
 
 # A thin wrapper to adapt a DynamicPPL.LogDensityFunction to the interface expected by
@@ -105,11 +105,9 @@ function DynamicPPL.marginalize(
     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
+        # The type assertion helps in cases where the model is type unstable and thus
+        # `varname_ranges` may have an abstract element type.
+        (ldf._varname_ranges[vn]::RangeAndLinked).range
     end
     mld = MarginalLogDensities.MarginalLogDensity(
         LogDensityFunctionWrapper(ldf, varinfo),

src/contexts/init.jl

@@ -206,13 +206,17 @@ an unlinked value.
 
 $(TYPEDFIELDS)
 """
-struct RangeAndLinked
+struct RangeAndLinked{T<:Tuple}
     # indices that the variable corresponds to in the vectorised parameter
     range::UnitRange{Int}
     # whether it's linked
     is_linked::Bool
+    # original size of the variable before vectorisation
+    original_size::T
 end
 
+Base.size(ral::RangeAndLinked) = ral.original_size
+
 """
     VectorWithRanges{Tlink}(
         varname_ranges::VarNamedTuple,
@@ -247,7 +251,12 @@ struct VectorWithRanges{Tlink,VNT<:VarNamedTuple,T<:AbstractVector{<:Real}}
 end
 
 function _get_range_and_linked(vr::VectorWithRanges, vn::VarName)
-    return vr.varname_ranges[vn]
+    # The type assertion does nothing if VectorWithRanges has concrete element types, as is
+    # the case for all type stable models. However, if the model is not type stable,
+    # vr.varname_ranges[vn] may infer to have type `Any`. In this case it is helpful to
+    # assert that it is a RangeAndLinked, because even though it remains non-concrete,
+    # it'll allow the compiler to infer the types of `range` and `is_linked`.
+    return vr.varname_ranges[vn]::RangeAndLinked
 end
 function init(
     ::Random.AbstractRNG,

src/logdensityfunction.jl

@@ -330,7 +330,10 @@ function get_ranges_and_linked_metadata(md::Metadata, start_offset::Int)
     for (vn, idx) in md.idcs
         is_linked = md.is_transformed[idx]
         range = md.ranges[idx] .+ (start_offset - 1)
-        all_ranges = BangBang.setindex!!(all_ranges, RangeAndLinked(range, is_linked), vn)
+        orig_size = varnamesize(vn)
+        all_ranges = BangBang.setindex!!(
+            all_ranges, RangeAndLinked(range, is_linked, orig_size), vn
+        )
         offset += length(range)
     end
     return all_ranges, offset
@@ -341,7 +344,10 @@ function get_ranges_and_linked_metadata(vnv::VarNamedVector, start_offset::Int)
     for (vn, idx) in vnv.varname_to_index
         is_linked = vnv.is_unconstrained[idx]
         range = vnv.ranges[idx] .+ (start_offset - 1)
-        all_ranges = BangBang.setindex!!(all_ranges, RangeAndLinked(range, is_linked), vn)
+        orig_size = varnamesize(vn)
+        all_ranges = BangBang.setindex!!(
+            all_ranges, RangeAndLinked(range, is_linked, orig_size), vn
+        )
         offset += length(range)
     end
     return all_ranges, offset

src/varname.jl

@@ -41,3 +41,28 @@ Possibly existing indices of `varname` are neglected.
 ) where {s,missings,_F,_a,_T}
     return s in missings
 end
+
+# TODO(mhauru) This should probably be Base.size(::VarName) in AbstractPPL.
+"""
+    varnamesize(vn::VarName)
+
+Return the size of the object referenced by this VarName.
+
+```jldoctest
+julia> varnamesize(@varname(a))
+()
+
+julia> varnamesize(@varname(b[1:3, 2]))
+(3,)
+
+julia> varnamesize(@varname(c.d[4].e[3, 2:5, 2, 1:4, 1]))
+(4, 4)
+"""
+function varnamesize(vn::VarName)
+    l = AbstractPPL._last(vn.optic)
+    if l isa Accessors.IndexLens
+        return reduce((x, y) -> tuple(x..., y...), map(size, l.indices))
+    else
+        return ()
+    end
+end