The ‘native’ data structure for a Turing model is VarNamedTuple, which is a mapping from VarNames to values. For example, rand(model) returns a VarNamedTuple:
usingTuringTuring.setprogress!(false)@modelfunctiondemo() x ~Normal(0, 1) y ~Normal(x, 1)endmodel =demo()vnt =rand(model)
[ Info: progress logging is disabled globally
VarNamedTuple
├─ x => -0.6127979577348269
└─ y => -1.5200751870750104
and you can likewise use a VarNamedTuple to specify initial parameters to Turing’s inference algorithms:
Now, VarNamedTuple is a rich data structure which faithfully retains names and types according to the structure of the model, which is why it is our default data structure.
However, many packages in the Julia ecosystem work with array inputs and outputs. Indeed this is even the case for many of Turing’s dependencies, such as AdvancedHMC, Optimization, and AdvancedVI: in these cases, Turing is responsible for performing the ‘translation’ between the VarNamedTuple and the vectorised format that these packages require.
This page describes how to convert a Turing model into a vectorised format such that it can be passed to these packages, and how to convert the outputs back into a VarNamedTuple format for analysis. This is useful if you need to carry out more advanced analysis with a Turing model with packages that Turing does not have wrappers for.
Introduction
Before we demonstrate any code, we will first explain conceptually what we need to actually accomplish this vectorisation. Given a VarNamedTuple of parameter values, we need two pieces of information to flatten it into a vector:
How to convert each individual value into a vector; and
How to concatenate those vectors together into a single vector.
We accomplish this by specifying a transform function for each parameter, as well as a range for each parameter. For example, conceptually, the VarNamedTuple above with x and y can be flattened with a representation that looks something like this.
(Please note that you don’t have to define any of this yourself! Turing has built-in functionality for this, which we’ll cover in a moment. This is just to illustrate the principle behind it.)
struct TransformPlusRange{F} transform::F range::UnitRange{Int}endtovec(f) = [f]info =@vntbegin x :=TransformPlusRange(tovec, 1:1) y :=TransformPlusRange(tovec, 2:2)end
VarNamedTuple
├─ x => TransformPlusRange{typeof(tovec)}(tovec, 1:1)
└─ y => TransformPlusRange{typeof(tovec)}(tovec, 2:2)
To flatten the parameters, we can then do the following:
# Begin by allocating a vector of the appropriate length.params =Vector{Float64}(undef, 2)# Iterate through the original `VarNamedTuple`. For each key:for (vn, value) inpairs(vnt) vn_info = info[vn]# Use the transform to vectorise the value vectorised_value = vn_info.transform(value)# Store it in the appropriate range of the allocated vector params[vn_info.range] = vectorised_valueendparams
Turing abstracts all of this away for you by providing a struct, LogDensityFunction, which wraps a model plus all of the necessary information above for vectorisation.
NoteMore details
In this page we will only cover a fairly high-level overview of the interface that LogDensityFunction exposes. In particular, we assume that you simply want to construct one of them, pass it to some package, and retrieve the results back in a VarNamedTuple format.
There are many more details of working with LogDensityFunction that can be useful if you are, for example, developing your own inference algorithms to work with Turing. We refer the interested reader to the DynamicPPL documentation for more information on this.
To construct a LogDensityFunction, you can simply pass a model:
@modelf() = x ~Dirichlet(ones(3))model =f()ldf =LogDensityFunction(model);
With this single-argument constructor, the vectorised representation of x will simply be itself (since Dirichlet is a multivariate distribution). Now, x also has the property that its elements must be non-negative and sum to 1. We can see this in action if we draw a randomised set of vectorised parameters:
By default, LogDensityFunction does not apply any extra transformations to the parameters, which is why the output above is still constrained to the simplex. However, such constraints can cause problems for some algorithms which expect to work in unconstrained space. To address this, you can instantiate a LogDensityFunction as follows:
LinkAll() specifies that we want all parameters in the model to be transformed to unconstrained vectors.
getlogjoint_internal says that the log-density we are interested in is the log-joint density, with any Jacobian adjustments for the transformations used above.
When we sample from this LogDensityFunction, we can see that the output is now unconstrained. In fact, the output is now a vector of length 2: since x must sum to 1, the third element of x is not needed since it can be recalculated from the first two elements.
LogDensityFunction is not just a way to perform parameter value conversion: it also allows you to compute the log probability density associated with these parameters.
In particular, the LogDensityFunction struct is designed to obey the LogDensityProblems.jl interface. In particular, you can use the following functions:
The vector passed to logdensity must have exactly dimension(ldf) elements. A vector of the wrong length throws an ArgumentError naming both lengths, rather than failing somewhere inside model evaluation or silently using only the first few elements.
Automatic differentiation
The LogDensityFunction that we created above does not yet know how to compute gradients of the parameters on its own. To enable integration with automatic differentiation (AD) packages, you can specify the adtype keyword argument:
When constructing a LogDensityFunction with an adtype, gradient preparation (through AbstractPPL’s prepare interface) will be performed automatically. This means that subsequent gradient calculations will be much faster, but also means that depending on the backend used, the construction of the LogDensityFunction may take a long time.
This allows you to compute not only logdensity, but also logdensity_and_gradient:
# We now have first-order AD capabilitiesLogDensityProblems.capabilities(ldf_ad)
If you need more information on how to choose an AD backend, please see the automatic differentiation page for more information.
Changing the log-density function
Most of the time, getlogjoint_internal is the appropriate kind of log-density function to use: it represents the log-joint density in the transformed space (‘internal’ is a slightly poor name, but it has stuck), and is what inference algorithms tend to expect.
However, there are also other options: for example LogDensityFunction(model, getloglikelihood, ...) will give you a log-density function that only computes the likelihood of the data given the parameters.
When instantiated, LogDensityFunction will automatically create its own mapping of VarNames to ranges and transforms (much like in the example above). If you are happy with the default transforms, then you would never need to worry about this. However, advanced users may wish to specify their own transforms.
In order to do this, you have to use a different constructor for LogDensityFunction, namely
Please see the DynamicPPL API documentation for more information on these arguments.
Minibatch log densities
Stochastic inference algorithms, such as those in AdvancedVI and stochastic-gradient variants of HMC, evaluate the log density on a minibatch of the data and scale the likelihood up to the full dataset. subsample builds a LogDensityFunction for one such minibatch.
The observation has to be written with independent_distribution, which declares that the data are conditionally independent draws, and the full dataset has to be supplied by conditioning rather than as a model argument. The model must have exactly one conditioned observation, and it must be the last probability-bearing statement. For example, placing an independent prior after that observation causes an error. Additional likelihood contributions through @addlogprob! are not supported.
subsample(model, indices, dataset_size) fixes the minibatch to the given indices, and subsample(model, batch_size, dataset_size) draws batch_size indices uniformly without replacement. The likelihood of the selected observations is multiplied by dataset_size / batch_size, while the prior enters once.
The minibatch stays fixed for the life of the LogDensityFunction, so an algorithm that wants a fresh batch at every step constructs a new one each time. Pass transform_strategy=LinkAll() when the algorithm needs unconstrained parameters. The DynamicPPL documentation on subsampling covers custom resamplers and the unbiasedness condition they must satisfy.
Conversions between VarNamedTuple and vectors
Finally, we will cover the functions that allow you to convert between VarNamedTuple and vector formats.
From VarNamedTuple to vector
We already saw that we can draw random samples from a LogDensityFunction in vector format:
By default, this draws from the prior of the model, using InitFromPrior (see Sampling options for more information on this). If you want vectorised samples for a specific set of parameters, you can provide an initialisation strategy that uses those parameters:
vnt =rand(model)
VarNamedTuple
└─ x => [0.1470513930385781, 0.6698148126017934, 0.18313379435962848]
(Note that this does technically mean that the output is no longer truly ‘random’.)
From vector to VarNamedTuple
Now suppose that you have run an inference algorithm using a LogDensityFunction and it has given you back a set of ‘best’ parameters (as a vector, of course):
This also helpfully computes the log-probabilities associated with this set of parameters (if you do not need this, you can disable it with ParamsWithStats(...; include_log_probs=false)). You can then access pws.params to get the VarNamedTuple of parameters. For example, to retrieve the original value of x:
These wrapper functions are designed for maximum ease of use. However, as mentioned above, sometimes you may want more fine-grained control. To do so, you should use the DynamicPPL API and in particular the DynamicPPL.init!! function. This is covered in much more detail in the DynamicPPL documentation!