Interface for transforms in AePPL

[rlouf]

TL;DR

I think we should remove the bijections from the extra_rewrites interface and move the automation to AeMCMC. This requires pushing for #78 or similar to allow conditioning on transformations of RandomVariables every time it is mathematically well-defined.

The interface for transforms is ambiguous

Let us consider the following model, where the $\mathcal{C}^+$ distribution is defined on $\mathbb{R}^+$:

import aesara
import aesara.tensor as at

srng = at.random.RandomStream(0)

mu_rv = srng.normal(0, 1)
sigma_rv = srng.halfcauchy(1)
Y_rv = srng.normal(mu_rv, sigma_rv)

Some samplers (for instance HMC) work better with variables that can take values in $\mathbb{R}$. AePPL proposes the following interface to allow to work in the transformed space with the above model:

import aeppl
from aeppl.transforms import TransformValuesRewrite, LogTransform

mu_vv = mu_rv.clone()
sigma_vv = sigma_rv.clone()
Y_vv = Y_rv.clone()

transforms_op = TransformValuesRewrite(
     {sigma_vv: LogTransform()}
)
logprob = aeppl.joint_logprob(
    {Y_rv: Y_vv, sigma_rv: sigma_vv, mu_rv: mu_vv},
    extra_rewrites=transforms_op
)

logprob_fn = aesara.function((Y_vv))

As discussed in #193 this interface is ambiguous: while sigma_vv is associated with sigma_rv with support $\mathbb{R}^+$, it now takes values in $\mathbb{R}$.

The interface for transforms is redundant

We could achieve the same result as the previous code snippet with the following:

import aeppl

mu_vv = mu_rv.clone()
Y_vv = Y_rv.clone()

sigma_tr = at.log(sigma_rv)
sigma_vv = sigma_tr.clone()

logprob = aeppl.joint_logprob({Y_rv: Y_vv, sigma_tr: sigma_vv, mu_rv: mu_vv})

The code is much shorter, the interface is not ambiguous. Why would we make callers remember the existence of TransformValuesRewrite and LogTransform when they already know about at.log? The current interface to apply these bijection transforms is not only confusing, it is completely redundant.

The above example currently does not work, for the reasons given in #119. This kind of conditioning should however be supported, independently from the above considerations. If anything, being able to clean the rewrite interface should give us extra motivation to move forward with e.g. #78.

DEFAULT_TRANSFORM is sampler-specific

AePPL does provide a utility for callers who do not know what transform should be applied; one would replace LogTransform with DEFAULT_TRANSFORM in the code above. However, this automation is sampler-specific and naturally belongs to AeMCMC.

Related to #142

Recommandations

I believe we should move the bijections out of the extra_rewrites interface, move the automation to AeMCMC and focus instead in AePPL on making conditioning on graphs with measurable transformations better, e.g. by working on #78.

Note that the only thing that would be removed is automation. The need to pass these rewrites to joint_logprob is already served in a much better way by defining transformed variables with Aesara Ops directly.

[rlouf]

Condition on a sequence of transforms of a single random variable

It is in theory straightforward to condition on a sequence of transfomations of a single RandomVariables, for instance in the following model:

import aesara
import aesara.tensor as at

srng = at.random.RandomStream(0)
x_rv = srng.normal(0, 1)
y = at.pow(at.log(x_rv), 2)

aesara.dprint(y)
# Elemwise{pow,no_inplace} [id A]
#  |Elemwise{log,no_inplace} [id B]
#  | |normal_rv{0, (0, 0), floatX, False}.1 [id C]
#  |   |RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7FE5D837A880>) [id D]
#  |   |TensorConstant{[]} [id E]
#  |   |TensorConstant{11} [id F]
#  |   |TensorConstant{0} [id G]
#  |   |TensorConstant{1} [id H]
#  |TensorConstant{2} [id I]

First, we need to provide x_rv a value by applying to y's value the successive reverse transformations. For that we simply need to walk between y and x_rv and apply the backward transformations successively.

y_vv = y.clone()
x_vv = at.exp(at.sqrt(y_vv))

To compute the probability $P(Y=y)$ one needs to follow the same path, accumulate the values of the log of the determinant of the jacobian of the successive transformation:

# P(y=y_vv)
logprob_y = log_det_jac_exp + log_det_jac_pow + _logprob(x_rv.op, x_vv)

Should we use Aesara Ops to represent transformations between measures?

Not entirely clear, since we'd need to register the log-jacobian determinant anyway + some transfomations between measures cannot be cast to Aesara Ops like ordering, sum-to-fixed-value constraint, etc.

x_rv = at.random.normal(0, 1)
y = at.exp(x_rv)

x_rv is not just a TensorVariable of type $\mathbb{R}$, it is a MeasurableVariable of type $\mathbb{M}\mathbb{R}$. When we're transforming a measurable variable we're transforming the underlying type as well as the measure, and Aesara Ops theoretically only operate on the underlying type.

Conditioning with ValuedVariables

If we have a similar mechanism to #78's ValuedVariables we would set the value of x_vv in the graph to at.exp(at.sqrt(y_vv)) so we can condition on the random variables that have x_rv as a parent in the model graph, thus solving #119.

Condition on the function of several random variables

We may want to condition on $P(Y=y)$ where $Y$ is the result of an algebraic combination of random variables:

import aeppl
import aesara
import aesara.tensor as at

srng = at.random.RandomStream(0)
s_rv = srng.normal(0, 1)
t_rv = srng.normal(0, 1)

y = s_rv / t_rv

y_vv = y.clone()
logprob = aeppl.joint_logprob({y: y_vv})

This situation is very different from the previous one: the distribution of $Y$ is called a ratio distribution and there are only a few cases in which we have a closed-form expression for this distribution, and thus only a few cases for which we can compute $P(Y=y)$. In this example, the ratio of $\operatorname{Normal}$ distributions has a know distribution:i $\operatorname{C}^+$. This knowledge is typically encoded in AeMCMC.

Using relations in AeMCMC and miniKanren to find closed-form expressions

To find those cases where a closed-form solution is available we should apply algebra-related rewrites from AeMCMC to the sub-graph that contains the variables and its RV parents beforehand; if such a closed-form is available we can then easily condition on y. This is an example of transform that has its place in AePPL since it does not change the nature of the variables over which we are computing (of factorizing) the joint probability distribution.

Furthermore, we should in this case use miniKanren to walk over equivalent representations of this subgraph until we've found a form for which we can compute the logprob. This would solve the concerns with the current implementation in #184. Performance should not be a concern here since we only need to work on one subgraph at a time.

Conditioning with ValuedVariables

Note that in this case, if we had ValuedVariables, we would set s_vv = y_vv * t_vv and t_vv = s_vv / y_vv before moving on with the other variables over which the joint probability is computed.

[brandonwillard]

Minor note: we might want to use at.exp(x_rv) in our examples instead, since it's better defined and corresponds to a log-normal transformation.

Also, we might need to highlight what it is specifically about miniKanren that's helpful in situations like this (e.g. its complete search, "relationally encoded" logic, etc.), in case others aren't familiar.