Refactor MeasurableElemwise and its use

[brandonwillard]

This PR refactors the design of MeasurableElemwise and the rewrites that use it (i.e. the changes mostly introduced by #26).

These changes make it easier for others to extend the coverage of measurable Elemwise transforms using basic rewrite DB registration. They also put each transformation rewrite under the management of Aesara's rewrite system independently, which means that they can be applied more efficiently (e.g. hash-based lookups on distinct Op instances) and their application can be enabled/disabled and tracked with more granularity—among other things.

The additions in #144 have also been added.

• Consider replacing the negation and subtraction rewrites with direct transform-constructing rewrites to avoid breaking canonicalization (see the comments in Implement measurable transforms for division, subtraction and negation #144)

Closes #170

[codecov]

Codecov Report

Base: 95.08% // Head: 95.15% // Increases project coverage by +0.06% 🎉

Coverage data is based on head (369b835) compared to base (d0c009d).
Patch coverage: 99.08% of modified lines in pull request are covered.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #208      +/-   ##
==========================================
+ Coverage   95.08%   95.15%   +0.06%     
==========================================
  Files          12       12              
  Lines        1974     2023      +49     
  Branches      258      253       -5     
==========================================
+ Hits         1877     1925      +48     
- Misses         55       56       +1     
  Partials       42       42              

Impacted Files	Coverage Δ
aeppl/cumsum.py	100.00% <ø> (ø)
aeppl/joint_logprob.py	95.00% <ø> (ø)
aeppl/mixture.py	96.59% <ø> (ø)
aeppl/tensor.py	85.84% <ø> (ø)
aeppl/transforms.py	96.45% <98.70%> (+0.20%)	⬆️
aeppl/abstract.py	100.00% <100.00%> (ø)
aeppl/censoring.py	96.22% <100.00%> (+0.22%)	⬆️
aeppl/rewriting.py	94.17% <100.00%> (ø)
aeppl/scan.py	94.73% <100.00%> (ø)

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.

[rlouf]

This looks great, and now that it's cleaned up it looks more obvious to me what can be further improved.

Do you want to also go ahead and include the changes in #184 or do you want to leave it for future work?

More general comment

I am not convinced that in the simple situations AePPL currently supports we need to rewrite the graph at all, although it could be useful to have a "canonicalization" phase where we could rewrite transforms of RandomVariable when they return a known distribution (ratio of StandardNormalRV for instance). Here, the model graphs contain almost all the information we need. To be able to compute the model's logdensity without rewriting the graph we would also need:

• To register every RandomVariable's logdensity.
• To register the inverse transform of Aesara operators to compute the upstream value variables;
• To register the way transforms act on logdensities.

However, pushed a little further, rewriting the graph could be very interesting for AeMCMC. We would need to replace every random and observed variable in the graph by their corresponding measure, including their domain and logdensity graph. That means starting from the known measures and pushing them through the transforms, thus getting a graph of Measures; transforms are then irrelevant if we also keep track of the relation between value variables in parallel. This is probably what we should be aiming for in these rewrites, instead of replacing Aesara operators with RVTransforms.

[brandonwillard]

Do you want to also go ahead and include the changes in #184 or do you want to leave it for future work?

Let's do that one in a follow-up.

More general comment

Yes, there's a lot to rework here. I'm going to merge this one so that—in the meantime—our existing functionality is preserved and better designed. At the very least, it should be a lot easier to rewrite/redesign in this form compared to the current.