Allow LogDensityProblems.logdensity to take a named tuple

[JaimeRZP]

Dear Team!

Some samplers such as MH want to provide LogDensityProblems.logdensity with a named tuple of parameters.
However, LogDensityProblems.logdensity only suports passing the parameters as a vector.
At the moment this is bypassed in Turing by passing the named tuple to the VarInfo, using DynamicPPL to evaluate the logp at that VarInfo and then resetting the VarInfo to its old state (see code below, also here:)

const MHLogDensityFunction{M<:Model,S<:Sampler{<:MH},V<:AbstractVarInfo} = Turing.LogDensityFunction{V,M,<:DynamicPPL.SamplingContext{<:S}}

function LogDensityProblems.logdensity(f::MHLogDensityFunction, x::NamedTuple)
    # TODO: Make this work with immutable `f.varinfo` too.
    sampler = DynamicPPL.getsampler(f)
    vi = f.varinfo

    x_old, lj_old = vi[sampler], getlogp(vi)
    set_namedtuple!(vi, x)
    vi_new = last(DynamicPPL.evaluate!!(f.model, vi, f.context))
    lj = getlogp(vi_new)

    # Reset old `vi`.
    setindex!!(vi, x_old, sampler)
    setlogp!!(vi, lj_old)

    return lj
end

Is there a better way of doing this?

All the best and thank you for your take,
Jaime

[yebai]

I’m not aware of any existing options. @devmotion might know?

[devmotion]

TransformVariables is the proto-typical examples of how to transform vectors (in unconstrained spaces) to NamedTuples of parameters in the (possibly) constrained space. It's also used in the LogDensityProblems docs: https://docs.julialang.org/en/v1/manual/types/#man-custom-pretty-printing

The interface of logdensity is supposed to be simple and predictable - and generally it is only guaranteed that a logdensity problem supports array-like inputs of the correct length. As in the docs, you can always make your logdensity problem a callable struct that accepts inputs of other types such as NamedTuples. To use this in a package you have to be certain though that the supported problems always accept the desired arguments.

[yebai]

I agree that it is better to keep logdensity interface predictable. If needed, one can define a convenience function like

Turing.jl/src/inference/Inference.jl

Lines 108 to 111 in ef27088

	const ESLogDensityFunction{M<:Model,S<:Sampler{<:ExternalSampler},V<:AbstractVarInfo} = Turing.LogDensityFunction{V,M,<:DynamicPPL.DefaultContext}
	function LogDensityProblems.logdensity(f::ESLogDensityFunction, x::NamedTuple)
	return DynamicPPL.logjoint(f.model, DynamicPPL.unflatten(f.varinfo, x))
	end