Sample remaining variables with the NUTS sampler

[rlouf]

In this PR we assign the NUTS sampler to the variables that have not been assigned a sampler in construct_sampler.

More details to come

Closes #47

[codecov]

Codecov Report

Base: 97.38% // Head: 97.41% // Increases project coverage by +0.02% 🎉

Coverage data is based on head (3a783e4) compared to base (ed60d94).
Patch coverage: 100.00% of modified lines in pull request are covered.

Additional details and impacted files

@@            Coverage Diff             @@
##             main      #68      +/-   ##
==========================================
+ Coverage   97.38%   97.41%   +0.02%     
==========================================
  Files           9        9              
  Lines         613      619       +6     
  Branches       60       58       -2     
==========================================
+ Hits          597      603       +6     
  Misses          5        5              
  Partials       11       11              

Impacted Files	Coverage Δ
aemcmc/basic.py	100.00% <100.00%> (ø)
aemcmc/nuts.py	98.03% <100.00%> (-0.04%)	⬇️

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]

Progress report

The NUTS sampler is integrated in construct_sampler; it uses the previous sampling_step to condition the logprob. There might be a performance trade-off with the current implementation as I re-initialize the state at every step; hopefully Aesara is able to identify the extra gradient computation is not needed, but that's something we will need to check.

It is currently tested on an example with NUTS-only. Doing so I discovered a bug when the variable is a scalar:

import aesara.tensor as at

srng = at.random.RandomStream(0)
tau_rv = srng.halfcauchy(0, 1, name="tau")
Y_rv = srng.halfcauchy(0, tau_rv, name="Y")

Which raises:

TypeError: Inconsistency in the inner graph of scan 'scan_fn' : an input and an output are associated with the same recurrent state and should have compatible types but have type 'TensorType(float64, (1,))' and 'TensorType(float64, (None,))' respectively.

This is reminiscent of shape issues I encountered in aehmc itself and I think the issue has to be solved there since it otherwise samples with no difficulty.

Most importantly, I need an example where there's both a Gibbs sampling step / closed form posterior and a NUTS sampling step.

To summarize, in order of priority:

• Add a test with mixed samplers
• Fix the shape issue for scalar random variables.
  This is due to RVs in Aesara having either a concrete or symbolic shape depending on the input. I opened an issue on Aesara: The shape of the RandomVariables is hard to reason about aesara#1252.
  Edit: I now specify the shape of inverse_mass_matrix when I initialize it. It doesn't solve the issue in Aesara, but this is good practice anyway and the tests pass.
• Assess whether we occur a performance penalty by initializing the state at each step
  NUTS graph are hard to navigate and we agreed this may be a good time to tackle Folding graph representation with Textual aesara#938

[rlouf]

This is ready for review:

• The shape issues were solved by specifying the shape of the mass matrix (but this does not solve the issue in Aesara);
• I moved the discussion about the general interface to Add information about sampling steps for the (human & machine) callers #71; this requires broader changes in the library. It probably requires more discussions and should be its own PR.
• Adaptation would build on these changes and is thus its own PR as well.
• I will open an issue to remember to check the graphs produced.

[brandonwillard]

Looks good; just some minor comments.