Decide how to deal with automatic broadcast in derived RVs

[ricardoV94]

Both #26 and #22 raise the question of how to deal with automatic RVs whose value_variable results from an intermediate broadcast of the base RV

loc = at.vector("loc")
y_base_rv = at.random.normal(0, 1, name="y_base_rv", size=2)
y_rv = loc + y_base_rv
y_val = y_rv.clone()

logp = joint_logprob(y_rv, {y_rv: y_val})
logp.eval({y_val: [0, 0, 0, 0], loc_val: [0, 0, 0, 0]})

The computed logp is only correct when loc is of size==2.
This might be even more complicated if the y_base_rv inputs are themselves variable (meaning it may itself have not a fixed size)

I see two ways we could deal with this:

1. Add a shape assert for y_base_rv.shape == y_val.shape in the returned logprob
2. Ignore the situation and put the onus on the users to specify correctly sized variables for derived logprobs
3. Something else?

[brandonwillard]

We defer to Aesara on these matters.

In other words, we let Aesara's error handling explain the issue to the user. It has special error messages for exactly these kinds of user-induced issues, so there's no reason to introduce redundancies to that effect in this library.

[brandonwillard]

Following up from internal discussions...

There appears to be two distinct—yet related—points that arise from this example:

1. size is being set to something specific, yet that specific value isn't necessarily honored by the resulting logp, and
2. y_rv = loc + y_base_rv involves a broadcasting operation when y_base_rv is, say, a scalar, and that has some interesting implications that might also not be honored by the resulting logp.

Point 1

The size situation can be illustrated with a simpler example:

Y_rv = at.random.normal(0, 1, name="Y", size=2)
y_vv = Y_rv.clone()

Y_logp = joint_logprob(Y_rv, {Y_rv: y_vv})

Y_logp.eval({y_vv: [0, 0]})
# array([-0.91893853, -0.91893853])

Everything seems to work fine, but we're also able to do the following:

Y_logp.eval({y_vv: [0, 0, 0]})
# array([-0.91893853, -0.91893853, -0.91893853])

In other words, the size=2 part of our model doesn't seem to matter.

What's really happening is that the "value" variable, y_vv, is being used to determine the general form of the log-probability/likelihood/density graph, and, because y_vv is just an Aesara symbolic vector, it's not constrained by the size=2 in the model-graph from which the log-probability was derived.

Currently, we don't use the size parameter to constrain the results of logprob, but we definitely could; however, there's a matter of interpretation regarding the output of joint_logprob that could make this too restrictive/unnecessary.

If we say that joint_logprob returns a graph that can be used to represent the given model-graph, then we don't need to force a more narrow kind of result that can only be used to represent the given model-graph.

Point 2

This point is a hypothetical, because we haven't implemented functionality for shifting and scaling just yet. It's quite easy to add such functionality for specific distributions right now (e.g. see the Algebraic Manipulations section here), but #26 should cover these cases more generally.

It seems like the original example is using size to say that Y_base_rv is a vector of i.i.d. normal variates. When these variates match loc, everything works out and Y_rv can be simplified to at.random.normal(loc, 1). For consistency, AePPL's rewrite logic will probably need to confirm that the shapes of loc and Y_base_rv conform without allowing a broadcast operation that changes the meaning of the model.

For instance, Y_base_rv = at.random.normal(0, 1, size=loc.shape) is one way of stating that Y_rv = at.random.normal(loc, 1) is always valid, but Y_base_rv = at.random.normal(0, 1, size=Y_size) where Y_size != loc.shape makes that rewrite generally invalid. These sorts of considerations are the responsibility of the rewrite/transformation logic.

Currently, we have a rewrite named naive_bcast_rv_lift that will naively lift a BroadcastTo through a RandomVariable node into its inputs—broadcasting those. This rewrite makes the overly strong assumption that the broadcasted variables should be i.i.d. variates that match the broadcasted shape, instead of a single scalar variate:

Y_rv = at.broadcast_to(at.random.normal(0, 1), (10,))
y_vv = Y_rv.clone()

Y_logp = joint_logprob(Y_rv, {Y_rv: y_vv})

Y_logp.eval({y_vv: [0, 1]})
# array([-0.91893853, -1.41893853])
Y_logp.eval({y_vv: [0, 1, 2]})
# array([-0.91893853, -1.41893853, -2.91893853])

This is somewhat confounded by the fact that y_vv—the given value variable—is necessarily a vector and not a scalar that represents the underlying normal variate. Given that type of value variable, one might be able to say that the requested joint_logprob graph is the i.i.d. interpretation of the model graph, while a scalar y_vv would actually represent the literal broadcasted interpretation.

I'll need to think about this for a minute longer, but, regardless, the value variable's type/shape is an important part of this no matter what, and especially since we don't even support the scalar y_vv case.

Also, we need to get back to naive_bcast_rv_lift and address the overly strong assumptions behind that rewrite. For instance, we can make it only work when one of the inputs/size is already consistent with the broadcasted shape. That will prevent it from turning a truly broadcasted random variable into something that's equivalent to i.i.d. random variates.