Value variables are not replaced upstream of unvalued RandomVariables

[ricardoV94]

import aesara.tensor as at
import aeppl

x_rv = at.random.normal(name="x")
y_rv = at.random.normal(x_rv, name="y")
z_rv = at.random.normal(y_rv, name="z")

x_vv = x_rv.clone()
z_vv = z_rv.clone()

# No value variable is given to `y_rv`
logp_dict = aeppl.factorized_joint_logprob({x_rv: x_vv, z_rv: z_vv})
logp_dict[z_vv].eval({x_vv: 0})  # MissingInputError

The logprob of z_rv does not depend on x_vv in the example above, because we short-circuit here:

aeppl/aeppl/utils.py

Lines 131 to 136 in 8ce3c63

	if (
	var.owner
	and (walk_past_rvs or not isinstance(var.owner.op, MeasurableVariable))
	and (var not in stop_at_vars)
	):
	new_vars.extend(reversed(var.owner.inputs))

[brandonwillard]

We should probably clarify the exact mathematical expressions we expect these quantities to reflect.

[ricardoV94]

That was brought up in #85

I assumed that for now we were working on the basis that RVs are allowed to remain unvalued in the returned graphs.

Otherwise my vote would be for those to be marginalized out (when possible).

[brandonwillard]

The logprob of z_rv does not depend on x_vv in the example above, because we short-circuit here:

aeppl/aeppl/utils.py

I never liked that conditional walking approach; it complicates things. It also implies that we're using the wrong graphs in some sense. If we can move away from it, then we'd probably be much better off.

[ricardoV94]

I assumed it was there for performance reasons? Otherwise I don't think any functionality depends on it at the moment.

[brandonwillard]

I assumed it was there for performance reasons? Otherwise I don't think any functionality depends on it at the moment.

No, it was just an easier way to do certain things in the beginning.

[rlouf]

Why do you expect the conditional distribution of z_rv to depend on x_vv here? It looks like the dependency path is blocked by the stochastic variable y_rv.

[ricardoV94]

Right, but y_rv loc should now be x_vv instead of x_rv.

[rlouf]

Right, but y_rv loc should now be x_vv instead of x_rv.

True, but it seems to be the case already. Here's what I get when I print the graph of joint_logprob:

import aesara
import aesara.tensor as at
import aeppl

x_rv = at.random.normal(name="x")
y_rv = at.random.normal(x_rv, name="y")
z_rv = at.random.normal(y_rv, name="z")

x_vv = x_rv.clone()
x_vv.name = "x_vv"
z_vv = z_rv.clone()
z_vv.name = "z_vv"

# No value variable is given to `y_rv`
logprob = aeppl.joint_logprob({x_rv: x_vv, z_rv: z_vv})
aesara.dprint(logprob)
# Sum{acc_dtype=float64} [id A]
#  |MakeVector{dtype='float64'} [id B]
#    |Sum{acc_dtype=float64} [id C]
#    | |Check{sigma > 0} [id D] 'x_vv_logprob'
#    |   |Elemwise{sub,no_inplace} [id E]
#    |   | |Elemwise{sub,no_inplace} [id F]
#    |   | | |Elemwise{mul,no_inplace} [id G]
#    |   | | | |TensorConstant{-0.5} [id H]
#    |   | | | |Elemwise{pow,no_inplace} [id I]
#    |   | | |   |Elemwise{true_div,no_inplace} [id J]
#    |   | | |   | |Elemwise{sub,no_inplace} [id K]
#    |   | | |   | | |x_vv [id L]
#    |   | | |   | | |TensorConstant{0.0} [id M]
#    |   | | |   | |TensorConstant{1.0} [id N]
#    |   | | |   |TensorConstant{2} [id O]
#    |   | | |Elemwise{log,no_inplace} [id P]
#    |   | |   |Elemwise{sqrt,no_inplace} [id Q]
#    |   | |     |TensorConstant{6.283185307179586} [id R]
#    |   | |Elemwise{log,no_inplace} [id S]
#    |   |   |TensorConstant{1.0} [id N]
#    |   |All [id T]
#    |     |Elemwise{gt,no_inplace} [id U]
#    |       |TensorConstant{1.0} [id N]
#    |       |TensorConstant{0.0} [id V]
#    |Sum{acc_dtype=float64} [id W]
#      |Check{sigma > 0} [id X] 'z_vv_logprob'
#        |Elemwise{sub,no_inplace} [id Y]
#        | |Elemwise{sub,no_inplace} [id Z]
#        | | |Elemwise{mul,no_inplace} [id BA]
#        | | | |TensorConstant{-0.5} [id BB]
#        | | | |Elemwise{pow,no_inplace} [id BC]
#        | | |   |Elemwise{true_div,no_inplace} [id BD]
#        | | |   | |Elemwise{sub,no_inplace} [id BE]
#        | | |   | | |z_vv [id BF]
#        | | |   | | |normal_rv{0, (0, 0), floatX, False}.1 [id BG] 'y'
#        | | |   | |   |RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7FEAF1AB3920>) [id BH]
#        | | |   | |   |TensorConstant{[]} [id BI]
#        | | |   | |   |TensorConstant{11} [id BJ]
#        | | |   | |   |x_vv [id L]
#        | | |   | |   |TensorConstant{1.0} [id N]
#        | | |   | |TensorConstant{1.0} [id N]
#        | | |   |TensorConstant{2} [id BK]
#        | | |Elemwise{log,no_inplace} [id BL]
#        | |   |Elemwise{sqrt,no_inplace} [id BM]
#        | |     |TensorConstant{6.283185307179586} [id BN]
#        | |Elemwise{log,no_inplace} [id BO]
#        |   |TensorConstant{1.0} [id N]
#        |All [id BP]
#          |Elemwise{gt,no_inplace} [id BQ]
#            |TensorConstant{1.0} [id N]
#            |TensorConstant{0.0} [id BR]

[ricardoV94]

@rlouf You're right. My example seems to always have worked because {x_rv: x_vv} is provided in the initial_replacements argument here:

aeppl/aeppl/joint_logprob.py

Line 143 in dead991

initial_replacements=replacements,

and that's all that is needed.

I thought it was an analogous issue to a similar function that we have in PyMC, but in that one we don't pass the initial_replacements.