Formalise VarInfo/AbstractVarInfo API

[yebai]

VarInfo currently implements and exports a large number of APIs, which sometimes can be redundant, confusing and hard to maintain. It is helpful to have a clear boundary what is internal to VarInfo and what can be reliably called by external modules such as sampling algorithms and the compiler itself.

[yebai]

Related #5 and #7.

[phipsgabler]

Cf. AbstractPPL discussion. I'm leaving this open for now for reference.

[yebai]

Some useful discussions

By @mohamed82008 #119 (comment)

Function redundancy is a big concern you may have if you have played around with varinfo.jl for a while. The main one I can think of is getval vs getindex and setval! vs setindex!. The way I think about it is that getval and setval! get and set the vectorized version of the value. So getval returns a vector and setval! expects a vector in even if the value was originally a scalar or matrix. getindex and setindex! on the other hand should reconstruct the shape of the original variable, and link/invlink the value based on whether the VarInfo is working in terms of the transformed or original values. Currently, setindex! doesn't do exactly that and but it was "fixed" in #115 to do exactly the above. So r = vi[vn] (or vi[vn, dist] in the PR) will always return a value in the domain of vn's distribution with the right shape and vi[vn] = r expects r to be in the domain of the distribution and in its natural shape.

Another redundancy that we have comes from the need for r = vi[vn] and vi[vn] = r but also r = vi[spl] and vi[spl] = r. The first 2 are used in a model call, while the other 2 are used in the step! function. These need to be defined for UntypedVarInfo as well as TypedVarInfo. But imo not much can be done about the need for those functions. VarInfo needs to define those functions one way or another because they are the main API of VarInfo.

By @devmotion #119 (comment)

IMO a major problem is that there is no clearly defined API right now (as discussed previously), which makes it difficult to actually come up with alternative AbstractVarInfo implementations (and, e.g., it's unclear if ThreadedVarInfo actually implements all required methods or maybe even too many). I think probably we should try to address this issue before adding any alternative implementations or rewrites of VarInfo.

To me it seems with the current version of VarInfo in mind a AbstractVarInfo object should at least support:

• getindex(varinfo, variables...) and setindex!(varinfo, value, variables...) for getting and setting values of variables
• copyto!(varinfo, vector), copyto!(vector, varinfo), and probably copyto!(varinfo, variables, vector) and copyto!(vector, varinfo, variables) for copying to/from vectors of variables
• getlogp, setlogp!, resetlogp!, and acclogp! for handling log probabilities

[yebai]

Also, by @phipsgabler in #283

The distinction between _getvalue, getval, and getindex is rather confusing. I suggested here to simplify the implementation of SimpleVarInfo as

haskey(vi::SimpleVarInfo, vn) = haskey(vi.θ, getsym(vn))

istrans(::SimpleVarInfo, vn::VarName) = false

getindex(vi::SimpleVarInfo, spl::SampleFromPrior) = vi.θ
getindex(vi::SimpleVarInfo, spl::SampleFromUniform) = vi.θ
# TODO: Should we do better?
getindex(vi::SimpleVarInfo, spl::Sampler) = vi.θ
function getindex(vi::SimpleVarInfo, vn::VarName{sym}) where {sym} 
    value = getproperty(nt, sym)
    return _getindex(value, vn.indexing)
end
getindex(vi::SimpleVarInfo, vns::AbstractArray{<:VarName}) = map(vn -> getindex(vi, vn), vns)

but that could lead to problems in the original VarInfo, which uses the other functions all over the place. Ideally, we get rid of _getvalue and getval completely and rely only on getindex.

(_getindex is a different case -- it performs the recursive indexing for VarNames, e.g. in x[1][2][3].)

EDIT: _getindex is already removed. 25 Jan 2022

[yebai]

Remove findvns: https://github.com/TuringLang/DynamicPPL.jl/blob/master/src/varinfo.jl#L1028

[yebai]

See also the design document in AbstractPPL.

[yebai]

Fixed by #417. See src/abstract_varinfo.jl