Implement measurable transforms for division, subtraction and negation

[ricardoV94]

Follow up to #26

This implements measurable rewrites for the following type of graphs:

x = at.random.gamma()
y = 1 / x
joint_logprob({y: y.type()})

As well as:

y = 5 / x
y = -x
y = 5 - x

It works by canonicalizing such operations to the form already understood by find_measurable_transforms (only a new condition for Reciprocals was added)

(5 / x) -> 5 * reciprocal(x)
(-x) -> x * (-1)
5 - x -> 5 + (x * (-1))

These canonicalizations are only applied when MeasurableVariables are involved to minimize disruption of the underlying graph

[codecov]

Codecov Report

Base: 94.82% // Head: 94.55% // Decreases project coverage by -0.27% ⚠️

Coverage data is based on head (df99f5b) compared to base (46f6775).
Patch coverage: 84.48% of modified lines in pull request are covered.

❗ Current head df99f5b differs from pull request most recent head 350a42e. Consider uploading reports for the commit 350a42e to get more accurate results

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #144      +/-   ##
==========================================
- Coverage   94.82%   94.55%   -0.28%     
==========================================
  Files          12       12              
  Lines        1933     1835      -98     
  Branches      289      276      -13     
==========================================
- Hits         1833     1735      -98     
- Misses         55       56       +1     
+ Partials       45       44       -1     

Impacted Files	Coverage Δ
aeppl/tensor.py	80.83% <82.60%> (-4.89%)	⬇️
aeppl/transforms.py	96.25% <91.66%> (-0.19%)	⬇️
aeppl/scan.py	94.73% <0.00%> (ø)
aeppl/abstract.py	100.00% <0.00%> (ø)
aeppl/rewriting.py
aeppl/censoring.py
aeppl/truncation.py	98.27% <0.00%> (ø)
aeppl/opt.py	94.00% <0.00%> (ø)
... and 3 more

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

☔ View full report at Codecov.
📢 Do you have feedback about the report comment? Let us know in this issue.

[brandonwillard]

Can you add a little more context to the description?

[ricardoV94]

Can you add a little more context to the description?

Done

[ricardoV94]

Added other rewrites that cover negation and subtraction of MeasurableVariables

[brandonwillard]

I'm very unsure about these fundamental (implicit) canonicalization changes made to rational expressions. Those could really create issues down the road, and they're something that should probably be used in Aesara as well, or not at all. In other words, if they're justified here, why aren't they justified for general use in Aesara?

Anyway, we might be able to try going forward with this—after renaming that class—but I still think it needs more consideration, especially since adding explicit canonicalizations isn't usually necessary (i.e. functionality-specific local rewrites can make the relevant rational form considerations instead).

[brandonwillard]

We need to start using the notation/approach mentioned here to clearly specify our rewrites at a high-level.

Specifically, the numerator not being equal to one could be compactly conveyed this way, along with the basic form of the rewrite itself.

[ricardoV94]

... forgot about this. Will do

[ricardoV94]

In other words, if they're justified here, why aren't they justified for general use in Aesara?

I think it makes sense to canonicalize in Aesara, except if weird float precision issues may crop up in general. We already have an issue opened for that, so if it goes through we can revert these ones here.

But even if we conclude they don't make sense in Aesara, here there are further context reasons to justify them. We know that in Aeppl most cases of unvalued RVs are meant to be replaced. If for some reason these canonicalizations are not desired, users can remove the rewrite from the standard database.

We also already have other canonicalizations that don't translate to measurable variables directly like the broadcast one.

The question is whether the trouble of affecting RV graphs that won't be replaced by value variables is larger than the cost of doing more refined intermediate rewrites that will be reverted automatically if unused...

[ricardoV94]

Alternatively, if we convert all unvalued RVs to MeasurableVariables before applying these canonicalizations, they will be undone when they are not needed by the same machinery in here:

aeppl/aeppl/opt.py

Line 321 in 394dc72

if rv_remapper.measurable_conversions:

[ricardoV94]

@brandonwillard I have not addressed the concerns about intermediate rewrites that simply canonicalize intermediate variables, besides restricting those rewrites to "unvalued measurable variables".

IMO until we figure out what we want to do with unmeasured RandomVariables in our graphs (as discussed in #174 and #85) we could keep assuming they are there to be replaced anyway, in which case these rewrites don't do any damage and make it simpler to cover a wider range of measurable Elemwise operations.

Anyway, I re-requested your review to check that this is the only concern left and discuss it.

[brandonwillard]

@brandonwillard I have not addressed the concerns about intermediate rewrites that simply canonicalize intermediate variables, besides restricting those rewrites to "unvalued measurable variables".

IMO until we figure out what we want to do with unmeasured RandomVariables in our graphs (as discussed in #174 and #85) we could keep assuming they are there to be replaced anyway, in which case these rewrites don't do any damage and make it simpler to cover a wider range of measurable Elemwise operations.

Anyway, I re-requested your review to check that this is the only concern left and discuss it.

We can always make these rewrites elective (i.e. a user must explicitly load/use them) for the time being; otherwise, adding rewrites that change the expectations of other rewrites (i.e. affect and/or determine canonical/normal forms) leads to serious downstream problems. For instance, if a / b is always converted to a * (1 / b) then subsequent rewrites expecting a / b (sub-)expressions are rendered useless. This could manifest as an inability to use certain basic Aesara rewrites with AePPL IR, or, at the very least, it could produce AePPL IR graphs that are inconsistent with the ones produced by Aesara for equivalent graphs.

Such "intermediate" rewrites can facilitate other rewrites that relate more directly to the desired functionality (i.e. the transforms in this PR), so the approach is very understandable. One important question that arises is these situations: "Can the intermediate rewrites be applied in isolation (i.e. not "committed" to the graph being rewritten) instead?".

In quite a few cases for which this isn't possible, I've found that the limiting factors involved an inability to perform reasonable graph comparisons (e.g. one would need to inefficiently walk both graphs). Core Aesara changes like aesara-devs/aesara#1110 and aesara-devs/aesara#1165 would address those issues.

In other cases, one needs intermediate rewrites to realize larger graph forms when, for example, there are no known sequences of "local" rewrites that will produce the desired forms consistently. Something like kanren can enumerate rewrite "trajectories" without making permanent changes to a graph, but that approach currently requires graph cloning (i.e. converting to etuples) and is only scalable when used in combination with some variety of memoization/tabling. We also don't have a consistent means of using Aesara rewrites in that context, so all the necessary rewrites need to be implemented in or converted to kanren goals. Core Aesara changes like aesara-devs/aesara#1082 and the use of Equality Saturation should allow us to address these cases more generally.

[ricardoV94]

Thanks for the reply.

I am not totally convinced why this matters here. Assuming we introduce other rewrites in the future that would try to match a = neg(x), I imagine they would also want to match b = mul(-1, x). Therefore such rewrite would need to 1) try to match both, in which case it doesn't matter that we changed from a to b or 2) save work and only match one, in which case it helps to have an "Aeppl canonicalized form" (b).

I understand why in Aesara we might not want to do that because of numerical stability. But here we are not dealing with direct computations yet, and we don't even know which, if any, IR form will lead to a more stable logprob (and the user much less), specially because this may be used as a building block for arbitrarily complex valued-variables.

Note that given the checks in place, the outputs of these intermediate rewrites will always be picked up by find_measurable_transforms and converted to a MeasurableVariable, so the final outcome would be completely equivalent if instead I expanded the find_measurable_transforms rewrite to be able to directly match different variations and skip the intermediate forms.

The only side-effect happens with unvalued-orphan RandomVariables which we haven't decided on a policy for yet.

[brandonwillard]

In general, it looks like the transform approach in #26 and this PR is rather buggy and acceptable changes to the rewrite orders cause numerous problems. After fixing something like #170 and running this, it becomes very apparent.

I'm going to close this and open a replacement PR that fixes the aforementioned, and the MeasurableElemwise approach as well.