Add a function that constructs samplers

[brandonwillard]

This PR provides an initial implementation of #3—and the related interface function mentioned in #26. It provides a general sampler-constructor, aemcmc.basic.construct_sampler, that returns a dict mapping RandomVariables to their sample steps.

The current approach uses a Feature called SamplerTracker to track a dict from RandomVariables to all their discovered sample steps—even when there's more than one potential sampler for the same RandomVariable. Sample steps are discovered by walking the graph with standard local rewriters that write their results to the dict in SamplerTracker. This allows us to maintain the original observation variable graphs in relation to every other un-observed variable (i.e. so we can see when a variable is in a particular hierarchical relationship with another variable, etc.)

In order to get around some DimShuffle annoyances during unification/pattern-matching, a SubsumingElemwise Op was added and is used to replace Elemwise(DimShuffle(x), ...) graphs with SubsumingElemwise(x, ...) graphs (i.e. ones that subsume the DimShuffles). Since SubsumingElemwise inherits from OpFromGraph, those nodes can be expanded later on to reproduce the original Elemwise + DimShuffle sub-graphs.

• [x] Work out the design of the main "processing loop" This loop needs to iterate over all the unobserved RandomVariables and construct samplers for them. While doing so, it needs to use references to the previously constructed samplers' outputs.
• [x] Do something about RandomVariable rewriting The issue here is that RandomVariables that are canonicalized are no longer the same RandomVariables that the user created, so we need a means of keeping a map between the two. Lifting Ops through RandomVariables is one of the main ways this issue shows up. N.B.: This is also one situation in which we could use complete relations (i.e. two-way rewrites).
• [x] Restrict the kinds of DimShuffles SubsumingElemwise will subsume (i.e. limit to only ones that add the appropriate broadcast dimensions). In cases where the original graph was a Elemwise(DimShuffle1(DimShuffle2(x)), ...) and the two DimShuffle* are merged, we will need to un-merge/expand them in order to use SubsumingElemwise.
• [x] Add initial value checks to tests (i.e. make sure that initial value variables correctly replace their corresponding RandomVariables).
• [x] Finish refactoring aemcmc.gibbs (e.g. remove combination samplers in favor of construct_sampler, create more local_optimizers, generalize local_optimizer construction, etc.)

[codecov[bot]]

Codecov Report

Merging #45 (65c7a37) into main (20611eb) will decrease coverage by 2.54%. The diff coverage is 97.59%.

:exclamation: Current head 65c7a37 differs from pull request most recent head 39ac1a5. Consider uploading reports for the commit 39ac1a5 to get more accurate results

@@            Coverage Diff             @@
##             main      #45      +/-   ##
==========================================
- Coverage   99.74%   97.20%   -2.55%     
==========================================
  Files           7        9       +2     
  Lines         391      572     +181     
  Branches       31       62      +31     
==========================================
+ Hits          390      556     +166     
- Misses          0        5       +5     
- Partials        1       11      +10     

Impacted Files	Coverage Δ
aemcmc/gibbs.py	91.87% <93.05%> (-8.13%)	:arrow_down:
aemcmc/opt.py	98.67% <98.67%> (ø)
aemcmc/basic.py	100.00% <100.00%> (ø)
aemcmc/conjugates.py	100.00% <100.00%> (ø)
aemcmc/dists.py	100.00% <100.00%> (ø)

---

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 20611eb...39ac1a5. Read the comment docs.

[brandonwillard]

We now have a complete working example in the test test_basic.py:test_create_gibbs.

[brandonwillard]

I've finished refactoring the sampler steps and filled out the docstings, so this should be ready to merge when/if it passes.