The behavior when conditioning on transforms of measures is incorrect

[rlouf]

The logprob term for z in the following graph depends on x_rv when it should depend on at.log(y_vv):

import aeppl
import aesara
import aesara.tensor as at

srng = at.random.RandomStream()
x = srng.normal(name="x")
y = at.exp(x)
z = srng.normal(x, name="z")

y_vv = y.clone()
z_vv = z.clone()
logprob, vvs = aeppl.joint_logprob(y, z)

aesara.dprint(logprob)
# Sum{acc_dtype=float64} [id A]
#  |MakeVector{dtype='float64'} [id B]
#    |Sum{acc_dtype=float64} [id C]
#    | |Elemwise{add,no_inplace} [id D]
#    |   |Check{sigma > 0} [id E]
#    |   | |Elemwise{sub,no_inplace} [id F]
#    |   | | |Elemwise{sub,no_inplace} [id G]
#    |   | | | |Elemwise{mul,no_inplace} [id H]
#    |   | | | | |TensorConstant{-0.5} [id I]
#    |   | | | | |Elemwise{pow,no_inplace} [id J]
#    |   | | | |   |Elemwise{true_div,no_inplace} [id K]
#    |   | | | |   | |Elemwise{sub,no_inplace} [id L]
#    |   | | | |   | | |Elemwise{log,no_inplace} [id M]
#    |   | | | |   | | | |<TensorType(float64, ())> [id N]
#    |   | | | |   | | |TensorConstant{0.0} [id O]
#    |   | | | |   | |TensorConstant{1.0} [id P]
#    |   | | | |   |TensorConstant{2} [id Q]
#    |   | | | |Elemwise{log,no_inplace} [id R]
#    |   | | |   |Elemwise{sqrt,no_inplace} [id S]
#    |   | | |     |TensorConstant{6.283185307179586} [id T]
#    |   | | |Elemwise{log,no_inplace} [id U]
#    |   | |   |TensorConstant{1.0} [id P]
#    |   | |All [id V]
#    |   |   |Elemwise{gt,no_inplace} [id W]
#    |   |     |TensorConstant{1.0} [id P]
#    |   |     |TensorConstant{0.0} [id X]
#    |   |Elemwise{neg,no_inplace} [id Y]
#    |     |Elemwise{log,no_inplace} [id Z]
#    |       |<TensorType(float64, ())> [id N]
#    |Sum{acc_dtype=float64} [id BA]
#      |Check{sigma > 0} [id BB] 'z_logprob'
#        |Elemwise{sub,no_inplace} [id BC]
#        | |Elemwise{sub,no_inplace} [id BD]
#        | | |Elemwise{mul,no_inplace} [id BE]
#        | | | |TensorConstant{-0.5} [id BF]
#        | | | |Elemwise{pow,no_inplace} [id BG]
#        | | |   |Elemwise{true_div,no_inplace} [id BH]
#        | | |   | |Elemwise{sub,no_inplace} [id BI]
#        | | |   | | |z [id BJ]
#        | | |   | | |normal_rv{0, (0, 0), floatX, False}.1 [id BK] 'x'
#        | | |   | |   |RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7F9003169C40>) [id BL]
#        | | |   | |   |TensorConstant{[]} [id BM]
#        | | |   | |   |TensorConstant{11} [id BN]
#        | | |   | |   |TensorConstant{0.0} [id O]
#        | | |   | |   |TensorConstant{1.0} [id P]
#        | | |   | |TensorConstant{1.0} [id P]
#        | | |   |TensorConstant{2} [id BO]
#        | | |Elemwise{log,no_inplace} [id BP]
#        | |   |Elemwise{sqrt,no_inplace} [id BQ]
#        | |     |TensorConstant{6.283185307179586} [id BR]
#        | |Elemwise{log,no_inplace} [id BS]
#        |   |TensorConstant{1.0} [id P]
#        |All [id BT]
#          |Elemwise{gt,no_inplace} [id BU]
#            |TensorConstant{1.0} [id P]
#            |TensorConstant{0.0} [id BV]

Example given originally in #119. We should be able to fix this after #78.

[rlouf]

Another example, which probably does not fail for the same reasons but illustrates how brittle the current framework is. The following returns, but the logprob is likely wrong as it does not include Y_vv:

import aeppl
import aesara.tensor as at

srng = at.random.RandomStream(0)

a_rv = srng.uniform(0, 1.0, name="A")
b_rv = srng.uniform(0, 1.0, name="B")
c_rv = srng.beta(.5, .5, name="C")
Y_rv = a_rv + c_rv * b_rv

logprob, (a_vv, b_vv, y_vv) = aeppl.joint_logprob(a_rv, b_rv, Y_rv)

But the following does not:

import aeppl
import aesara.tensor as at

srng = at.random.RandomStream(0)

a_rv = srng.uniform(0, 1.0, name="A")
b_rv = srng.uniform(0, 1.0, name="B")
c_rv = srng.beta(.5, .5, name="C")
Y_rv = a_rv + (a_rv - c_rv) * b_rv  # <--- this changed

logprob, (a_vv, b_vv, y_vv) = aeppl.joint_logprob(a_rv, b_rv, Y_rv)
# RuntimeError: The logprob terms of the following random variables could not be derived: {Elemwise{add,no_inplace}.0}

I found this out while trying to sample from the PERT distribution (aesara-devs/aemcmc#88 (reply in thread)) defined as a transformed Beta distribution:

import aemcmc
import aesara.tensor as at
import numpy as np

srng = at.random.RandomStream(0)

a_rv = srng.uniform(0, 1.0, name="A")
b_rv = srng.uniform(0, 1.0, name="B")
c_rv = srng.uniform(0, 1.0, name="C")

# Create some "observations"
rng = np.random.default_rng(2302)
y_obs = at.as_tensor(rng.uniform(0, 2, size=(10,)))

alpha = 1 + 4 * (b_rv - a_rv) / (c_rv - a_rv)
beta = 1 + 4 * (c_rv - b_rv) / (c_rv - a_rv)
X_rv = srng.beta(alpha, beta, size=(10,))
Y_rv = a_rv + (b_rv - a_rv) * X_rv

sampler, initial_values = aemcmc.construct_sampler({Y_rv: y_obs}, srng)
print(sampler.stages)

We should however be able to routinely do this if we don't want to write a custom logprob/associate custom transforms every time we want to add a new distribution in AePPL. We should be able to express most distributions as a graph that transforms other random variables, and condition on the output of this graph.

[brandonwillard]

But the following does not:

import aeppl
import aesara.tensor as at

srng = at.random.RandomStream(0)

a_rv = srng.uniform(0, 1.0, name="A")
b_rv = srng.uniform(0, 1.0, name="B")
c_rv = srng.beta(.5, .5, name="C")
Y_rv = a_rv + (a_rv - c_rv) * b_rv  # <--- this changed

logprob, (a_vv, b_vv, y_vv) = aeppl.joint_logprob(a_rv, b_rv, Y_rv)
# RuntimeError: The logprob terms of the following random variables could not be derived: {Elemwise{add,no_i

My first thought is that the - c_rv isn't supported. Perhaps that's related to MeasurableTransform.valid_scalar_types. (See also #206.)

I'll take a look shortly.

[brandonwillard]

The following returns, but the logprob is likely wrong as it does not include Y_vv:

import aeppl
import aesara.tensor as at

srng = at.random.RandomStream(0)

a_rv = srng.uniform(0, 1.0, name="A")
b_rv = srng.uniform(0, 1.0, name="B")
c_rv = srng.beta(.5, .5, name="C")
Y_rv = a_rv + c_rv * b_rv

logprob, (a_vv, b_vv, y_vv) = aeppl.joint_logprob(a_rv, b_rv, Y_rv)

The log-density graph appears to have the returned value variable, y_vv, for Y_rv in it, so it might be fine.

[rlouf]

I updated my original example to reflect the current API, and it still fails. The other example fails for a different reason so I masked the comments and opened a separate issue.

My guess is that we need to add an extra step in the IR rewriting, where we walk from the variables we condition on starting with the ones that are not RandomVariables, here y. We set the value variable of the upstream RandomVariable with the transform of y_vv.

[brandonwillard]

Here's the AePPL IR produced in that example:

import aesara
import aesara.tensor as at

srng = at.random.RandomStream()
X_rv = srng.normal(name="X")
Y_rv = at.exp(X_rv)
Z_rv = srng.normal(X_rv, name="Z")


from aeppl.rewriting import construct_ir_fgraph

y_vv = Y_rv.clone()
y_vv.name = "y"
z_vv = Z_rv.clone()
z_vv.name = "z"
fgraph, _, _ = construct_ir_fgraph({Y_rv: y_vv, Z_rv: z_vv})


aesara.dprint(fgraph)
# ValuedVariable [id A] 2
#  |MeasurableElemwiseTransform{exp} [id B] 1
#  | |normal_rv{0, (0, 0), floatX, False}.1 [id C] 'X' 0
#  |   |RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7FB3D9071820>) [id D]
#  |   |TensorConstant{[]} [id E]
#  |   |TensorConstant{11} [id F]
#  |   |TensorConstant{0.0} [id G]
#  |   |TensorConstant{1.0} [id H]
#  |y [id I]
# ValuedVariable [id J] 5
#  |normal_rv{0, (0, 0), floatX, False}.1 [id K] 'Z' 4
#  | |RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7FB3D90713C0>) [id L]
#  | |TensorConstant{[]} [id E]
#  | |TensorConstant{11} [id F]
#  | |normal_rv{0, (0, 0), floatX, False}.1 [id M] 'X' 3
#  | | |RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7FB3D9071820>) [id D]
#  | | |TensorConstant{[]} [id E]
#  | | |TensorConstant{11} [id F]
#  | | |TensorConstant{0.0} [id G]
#  | | |TensorConstant{1.0} [id H]
#  | |TensorConstant{1.0} [id H]
#  |z [id N]