Allow Minibatch of derived RVs and deprecate generators as data #7480
+159
−184
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
ricardoV94
added
enhancements
major
ricardoV94
commented
Aug 27, 2024
| mA, mB = pm.Minibatch(A, B, batch_size=10) | ||
|
|
||
| [draw_mA, draw_mB] = pm.draw([mA, mB]) | ||
| assert draw_mA.shape == (10,) | ||
| np.testing.assert_allclose(draw_mA, draw_mB) | ||
|
|
||
| # Check invalid dims |
There was a problem hiding this comment.
This was already checked in the test above
ricardoV94
force-pushed
the
minibatch_censored
branch
2 times, most recently
from
479c4a4 to
290a643
Compare
ricardoV94
force-pushed
the
minibatch_censored
branch
from
290a643 to
49542b5
Compare
Codecov ReportAttention: Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #7480 +/- ##
==========================================
- Coverage 92.16% 92.15% -0.02%
==========================================
Files 103 103
Lines 17214 17224 +10
==========================================
+ Hits 15866 15873 +7
- Misses 1348 1351 +3
|
1 task
ferrine
approved these changes
Sep 7, 2024
| mb_tensors = [tensor[mb_indices] for tensor in tensors] | ||
|
|
||
| # Wrap graph in OFG so it's easily identifiable and not rewritten accidentally | ||
| *mb_tensors, _ = MinibatchOp([*tensors, rng], [*mb_tensors, rng_update])(*tensors, rng) |
| @@ -666,6 +672,9 @@ class GeneratorOp(Op): | |||
| __props__ = ("generator",) | |||
|
|
|||
| def __init__(self, gen, default=None): | |||
| warnings.warn( | |||
| mA, mB = pm.Minibatch(A, B, batch_size=10) | ||
|
|
||
| [draw_mA, draw_mB] = pm.draw([mA, mB]) | ||
| assert draw_mA.shape == (10,) | ||
| np.testing.assert_allclose(draw_mA, draw_mB) | ||
|
|
||
| # Check invalid dims |
| @@ -162,22 +159,7 @@ def fit_kwargs(inference, use_minibatch): | |||
|
|
|||
|
|
|||
| def test_fit_oo(inference, fit_kwargs, simple_model_data): | |||
| # Minibatch data can't be extracted into the `observed_data` group in the final InferenceData | |||
There was a problem hiding this comment.
Nope I allow to extract the data in the idata, it extracts the whole data
| minibatch_idx = minibatch_index(0, 10, size=(9,)) | ||
| AD_mt = AD[minibatch_idx] | ||
| TD_mt = TD[minibatch_idx] | ||
| AD_mt, TD_mt = Minibatch(AD, TD, batch_size=9) |
|
I've created an issue to continue this work later and improve scalability of minibatches #7496 |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Description
This PR fixes issues related to Minibatch indexing reported in https://discourse.pymc.io/t/warning-using-minibatch-and-censored-together-rng-variable-has-shared-clients/14943 and extends the MinibatchRV functionality for derived RVs.
Minibatch value variables are uniquely tricky because they are random graphs, that can share RNG with other variables in the forward / logp graph. As such we need to make sure they are not mutated for the default updates to work. We tried some tricks in the past but as revealed in the discourse issue that was not enough. This PR solves the problem by encapsulating the random graph in an OpFromGraph so that the inner graph will not be touched by PyMC logprob derivation routines. It will still be inlined in the final compiled functions to avoid overhead.
I also decided to deprecate Generators as data, which showed up in some of the refactoring. The GeneratorOp is not a true Op, which should not have any side-effects. It is also not compatible with non default backends like Numba and JAX that we are moving towards to. If needed, the logic should be handled by the sampler by consuming the generator and setting the values before subsequent function calls.