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]

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]