Add a function to create a tensor that represents the loop over the posterior #7885
Conversation
lucianopaz
changed the title
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #7885 +/- ##
=======================================
Coverage 92.95% 92.96%
=======================================
Files 116 116
Lines 18863 18904 +41
=======================================
+ Hits 17534 17574 +40
- Misses 1329 1330 +1
🚀 New features to boost your workflow:
|
|
I would rather work on pytensor to refactor AdvancedSubtensor. The problem is described in, pymc-devs/pytensor#541 And/or allow |
| samples = clone_while_sharing_some_variables( | ||
| outputs, replace=replace, kept_variables=non_sequences | ||
| ) |
There was a problem hiding this comment.
I know VI does something like this, either by passing things that shouldn't be replaced as keys:values, or by using graph_replace, can you check if the same approach works?
There was a problem hiding this comment.
| @@ -1083,3 +1090,122 @@ def vectorize_over_posterior( | |||
| f"The following random variables found in the extracted graph: {remaining_rvs}" | |||
| ) | |||
| return vectorized_outputs | |||
|
|
|||
|
|
|||
| def loop_over_posterior( | |||
There was a problem hiding this comment.
instead of a separate function have an argument: use_scan: bool = False in the previous function? Most logic should be the same.
This flatten batch -> scan -> reshape logic should go in PyTensor, it's needed in general
Description
This PR adds the
loop_over_posteriorfunction. It complementsvectorize_over_posteriorand can be used whenevervectorize_graphis not implemented for some nodes in the computational graph (i.e. some kinds ofAdvancedSetSubtensorraiseNotImplementedErrorwhen trying to vectorize them).Related Issue
Checklist
Type of change
📚 Documentation preview 📚: https://pymc--7885.org.readthedocs.build/en/7885/