Tilde-statements in-depth

The Model evaluation page gives a high-level overview of how to control model evaluation in DynamicPPL. This page goes into more detail about exactly how initialisation strategies, transform strategies, and accumulators are implemented, and how they combine to give the overall semantics described on that page.

The @model macro

Each tilde-statement, say, x ~ Normal(), is transformed by the @model macro into something like the following pseudocode (the functions like is_fixed do not have those exact names in DynamicPPL but are conceptually equivalent). The effective binding determines the LHS variable's role; see Binding rules. Arguments on the LHS supply observations by default. Every argument binding replaces the argument before the body runs, so an observed argument LHS variable reads its current body value; a local LHS variable reads its binding at the tilde. Fixed argument LHS variables reset to the bound value at the tilde and require static shape. If you're interested in the gory details you can run @macroexpand @model f() = x ~ Normal() in the REPL.

begin
    vn = @varname(x)
    dist = Normal()
    # To understand the need for `template`, see the VarNamedTuple docs.
    template = x

    if is_fixed(__model__, vn)
        raw_x = get_fixed_value(__model__, vn)

    elseif is_conditioned(__model__, vn)
        # Conceptual branch: the compiler knows whether x is an argument.
        conditioned_x = is_argument(x) ? x : get_conditioned_value(__model__, vn)
        raw_x, __varinfo__ = tilde_observe!!(
            __prefix__, __prefix_template__, dist, conditioned_x, vn, template, __varinfo__
        )

    else
        raw_x, __varinfo__ = tilde_assume!!(
            __model__, __context__, dist, vn, template, __varinfo__
        )
    end

    x = raw_x
end

Here, __prefix__ and __prefix_template__ are the model's prefix and its storage template, passed directly to the exported tilde_observe!! entry point.

We won't go into detail about every part of this code; by far the most interesting part is the call to tilde_assume!!. Every latent tilde-statement vn ~ dist is transformed into one such call.

As described on the Model evaluation page, there are three stages to a latent tilde-statement:

1. Initialisation: get an TransformedValue from the initialisation strategy.
2. Transformation: figure out the untransformed (raw) value and the transformed value (where necessary); compute the relevant log-Jacobian.
3. Accumulation: pass all the relevant information to the accumulators, which individually decide what to do with it.

The method for tilde_assume!! (with InitContext) more or less implements this logic directly with three lines of code. The implementation in src/contexts/init.jl follows this structure:

function DynamicPPL.tilde_assume!!(ctx::InitContext, dist, vn, template, vi)
    # 1. Initialisation
    init_tval = DynamicPPL.init(ctx.rng, vn, dist, ctx.strategy)

    # 2. Transformation
    x, tval, logjac = apply_transform_strategy(ctx.transform_strategy, init_tval, vn, dist)

    # 3. Accumulation
    vi = DynamicPPL.accumulate_assume!!(vi, x, tval, logjac, vn, dist, template)
    return x, vi
end

For tilde_observe!!, the code is very similar, but even easier: the value can be read directly from the data provided to the model, so there is no need for an initialisation step. Since the value is already untransformed, we can skip the second step. Finally, accumulators must behave differently: e.g. incrementing the likelihood instead of the prior. That is accomplished by calling accumulate_observe!! instead of accumulate_assume!!. Neither named nor literal observations dispatch on the evaluation context. Fixed LHS variables bypass accumulation, and tracked assignments (:=) write directly to the raw-value accumulator.

In the following sections, we stick to the three sections of tilde_assume!!.

Customise value selection by implementing an initialisation strategy, not by introducing another context type.

Initialisation

init_tval = DynamicPPL.init(ctx.rng, vn, dist, ctx.strategy)

The initialisation step is handled by the init function, which dispatches on the initialisation strategy. For example, if ctx.strategy is InitFromPrior(), then init() samples a value from the distribution dist.

As discussed in the Initialisation strategies page, this step, in general, does not return just the raw value (like rand(dist)). It returns an DynamicPPL.TransformedValue, which represents a value that may have been transformed. In the case of InitFromPrior(), the value is of course not transformed; we return a DynamicPPL.TransformedValue(..., NoTransform()) wrapping the sampled value.

However, consider the case where we are using parameters stored inside a VarInfo: the value may have been stored either as a vectorised form, or as a linked vectorised form. In this case, init() will return a TransformedValue wrapping the vectorised value, whose transform is either Unlink() (i.e., vectorised but not linked), or DynamicLink() (i.e., vectorised and linked).

The reason why we return this wrapped value is because we want to avoid having to perform transformations multiple times. Each step is responsible for only performing the transformations it needs to. At this stage, there has not yet been any need for the raw value, so we do not perform any transformations yet. Thus, the TransformedValue is passed straight through and is used by the transformation step.

Transformation

x, tval, logjac = apply_transform_strategy(ctx.transform_strategy, init_tval, vn, dist)

There are three return values in this step, and they correspond to the three things that this step needs to do. They are all interconnected, which is why they are computed together inside apply_transform_strategy(): by doing so we can ensure that with_logabsdet_jacobian is only called a maximum of once per tilde-statement.

1. Get the raw (untransformed) value x

   At some point, we do need to perform the transformation to get the actual raw value. This is because DynamicPPL promises in the model that the variables on the left-hand side of the tilde are actual raw values.

   @model function f()
       x ~ dist
       # Here, `x` _must_ be the actual raw value.
       @show x
   end

   Thus, regardless of what we are accumulating, we will have to unwrap the transformed value provided by init().

2. Get the (possibly transformed) value tval

   In addition to the raw value, if the transform strategy indicates that we should treat vn as being in transformed space, we also need to compute the transformed value. This is because some accumulators may need to work with the transformed value instead of the raw value.

3. Compute the log-Jacobian logjac

   logjac is accumulated according to the transform specified by the transform strategy. The convention in DynamicPPL is that the log-Jacobian is always computed with respect to the forward transformation.

It is worth emphasising that whether a value is transformed or not is determined by the transform strategy provided to the model (i.e., ctx.transform_strategy), not the initialisation strategy (ctx.strategy). The reason for this is to allow a separation between the source of the values (initialisation) and how those values are to be interpreted (transform strategy).

This allows us to, for example, generate values from the (unlinked) prior but also calculate their log-density in transformed space and accumulate transformed values by combining InitFromPrior() with LinkAll(). It also allows us to read values from an existing VarInfo but interpret them as being in a different space by combining InitFromParams() with a different transform strategy: this corresponds exactly to the act of 'linking' a VarInfo.

Accumulation

vi = DynamicPPL.accumulate_assume!!(vi, x, tval, logjac, vn, dist, template)

accumulate_assume!!(vi::AbstractVarInfo, ...) will loop over all accumulators stored inside vi, and call each of their individual accumulate_assume!! methods. This method is responsible for deciding how to combine the information provided.

Accumulators are described in much more detail on the Accumulators page; please read that for more information!