Do not replace valued RVs in naive_bcast_rv_lift

[ricardoV94]

Closes #115

[brandonwillard]

It seems like this is trying to prevent rewrites on valued/observed random variables, but we definitely don't want to do that.

If relationships between rewritten terms aren't being preserved properly, then it sounds like there's an issue with the rewriting process.

[ricardoV94]

The more I think about this rewrite, the less convinced I am it ever makes sense. If you have x_rv and at_broadcast(x_rv, shape) as the outputs of a function, the actual values may become almost independent (subject to intricacies of the rng) once this rewrite is applied to the second variable.

For me, this only really makes sense to facilitate uncovering rv-valued variables.

Edit: The original bug here concerned the logprob graph not the random one, but the principle is almost the same. The new returned variable has no relationship with the original one, which is still part of the FunctionGraph.

This feels very specific to RandomVariables which are sort of an "input" variable and I don't think would be a problem with the other type of rewrites we are doing (although we have been limiting it to nodes that have corresponding value variables, so maybe it is a more general issue)

[brandonwillard]

The more I think about this rewrite, the less convinced I am it ever makes sense. If you have x_rv and at_broadcast(x_rv, shape) as the outputs of a function, the actual values may become almost independent (subject to intricacies of the rng) once this rewrite is applied to the second variable.

For me, this only really makes sense to facilitate uncovering rv-valued variables.

It sounds like you're referring to the representation issue addressed in #78.

The RandomVariables produced after lifting the broadcast operations aren't "proper" random variables. The entire purpose of this rewriting is to produce a log-probability graph, and those rewritten RandomVariables only serve as type-based measure information that we use in order to determine/derive the log-probabilities.

Simply put, when a BroadcastTo is lifted through a RandomVariable, we're not trying to produce the equivalent RandomVariable; we're constructing a RandomVariable of the same type, and with properly broadcastated parameters, that when evaluated at a given point will produce the same log-probability that the un-lifted RandomVariable would.

The graphs that make up our IR are never meant to be sampled, so RNG-based (in)dependence isn't relevant. For that matter, RNGs themselves aren't relevant to these graphs. If you're thinking of a rewrite that's valid in sample-space graphs, then you're not considering the space we're actually working within.

In other words, while two random variables may be distinct, their measures can still be equivalent, and that's what we're dealing with here. The confusion mostly arises from us not making/having a very formal distinction between our graphical representations of the two scenarios.

#78 addresses some of this by making such rewrites applicable only to distinct valued RandomVariable objects. We could also create replacement Ops for RandomVariables that—say—remove the RNG components and have a more revealing name (e.g. something having to do with them providing measure information), but I don't know how useful that would really be.

[brandonwillard]

This feels very specific to RandomVariables which are sort of an "input" variable and I don't think would be a problem with the other type of rewrites we are doing (although we have been limiting it to nodes that have corresponding value variables, so maybe it is a more general issue)

This sounds like a line of investigation that will lead to the real problem.

[brandonwillard]

If you rebase this onto main, the tests should run.

[ricardoV94]

This sounds like a line of investigation that will lead to the real problem.

I understand the real problem. I am not sure what other type of solution there is...

This change is not affecting our ability to parse logprob terms because you couldn't ever measure both the broadcasted node and the base rv at the same time anyway.

If you are trying to 1) measure the broadcasted node it will still work, and if you are trying to 2) measure the base RV node and only then broadcast the value this won't mess up the graph dependencies. In both cases the behavior is what we want for the IR.

Third case 3) is when you have several RandomVariables that you want to keep in the logprob graph and you have clients for both the base and broadcasted nodes. That would hit the "random graph" issue I mentioned above. The solution for that would be to limit this rewrite even more than what I did here and only allow it for 1) valued broadcasted nodes.

Of course if we did that (which this PR does not) we would limit the ability to measure 4) nested broadcasted nodes, which is a more general issue we still don't yet have a solution for with other rewrites.

[codecov]

Codecov Report

Merging #116 (170f769) into main (2019114) will increase coverage by 0.20%.
The diff coverage is 66.66%.

@@            Coverage Diff             @@
##             main     #116      +/-   ##
==========================================
+ Coverage   94.04%   94.25%   +0.20%     
==========================================
  Files          11       11              
  Lines        1529     1532       +3     
  Branches      217      218       +1     
==========================================
+ Hits         1438     1444       +6     
+ Misses         53       49       -4     
- Partials       38       39       +1     

Impacted Files	Coverage Δ
aeppl/opt.py	81.41% <66.66%> (+3.23%)	⬆️

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 04ae9c9...170f769. Read the comment docs.

[brandonwillard]

What if you added something like the following to naive_bcast_rv_lift:

    rv_map_feature = getattr(fgraph, "preserve_rv_mappings", None)

    if rv_map_feature is not None and rv_var in rv_map_feature.rv_values:
        val_var = rv_map_feature.rv_values[rv_var]
        val_var = at.broadcast_to(val_var, tuple(bcast_shape))
        rv_map_feature.update_rv_maps(rv_var, val_var, bcasted_node.outputs[1])

The main problem I see here is that naive_bcast_rv_lift simply isn't value-variable conscientious like the other rewrites are.