Refactor MiniBatch and stop using deprecated MRG sampler #6304

ferrine · 2022-11-15T07:54:23Z

fix #4523, suspend #6273
closes #6277

@ricardoV94, I've figured out the root cause of the seeding bug

ricardoV94 · 2022-11-15T08:14:21Z

@ferrine this does not fix #4523, VI is still using the MRG RandomStream. This can't be switched yet because of #6277

pymc/pymc/aesaraf.py

Line 52 in 5da2617

from aesara.sandbox.rng_mrg import MRG_RandomStream as RandomStream

pymc/pymc/aesaraf.py

Lines 872 to 895 in 5da2617

    
           _at_rng = RandomStream() 
        
           def at_rng(random_seed=None): 
        
               """ 
        
               Get the package-level random number generator or new with specified seed. 
        
               Parameters 
        
               ---------- 
        
               random_seed: int 
        
                   If not None 
        
                   returns *new* aesara random generator without replacing package global one 
        
               Returns 
        
               ------- 
        
               `aesara.tensor.random.utils.RandomStream` instance 
        
                   `aesara.tensor.random.utils.RandomStream` 
        
                   instance passed to the most recent call of `set_at_rng` 
        
               """ 
        
               if random_seed is None: 
        
                   return _at_rng 
        
               else: 
        
                   ret = RandomStream(random_seed) 
        
                   return ret

ricardoV94 · 2022-11-15T08:17:18Z

It does fix some lack of seeding, for which I think we don't have an issue open

pymc/variational/opvi.py

codecov · 2022-11-15T08:26:10Z

Codecov Report

Merging #6304 (1218d0c) into main (5692fc0) will increase coverage by 0.07%.
The diff coverage is 97.81%.

❗ Current head 1218d0c differs from pull request most recent head 82dfbf6. Consider uploading reports for the commit 82dfbf6 to get more accurate results

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #6304      +/-   ##
==========================================
+ Coverage   94.71%   94.79%   +0.07%     
==========================================
  Files         132      148      +16     
  Lines       26771    27488     +717     
==========================================
+ Hits        25357    26058     +701     
- Misses       1414     1430      +16

Impacted Files	Coverage Δ
pymc/distributions/__init__.py	`100.00% <ø> (ø)`
pymc/logprob/scan.py	`97.35% <ø> (ø)`
pymc/variational/approximations.py	`87.72% <50.00%> (+0.05%)`	⬆️
pymc/variational/opvi.py	`87.10% <85.00%> (+0.02%)`	⬆️
pymc/logprob/transforms.py	`97.74% <93.10%> (-1.59%)`	⬇️
pymc/distributions/distribution.py	`97.46% <98.23%> (+2.28%)`	⬆️
pymc/tests/distributions/test_distribution.py	`98.69% <99.24%> (+0.48%)`	⬆️
pymc/backends/arviz.py	`96.37% <100.00%> (+5.65%)`	⬆️
pymc/data.py	`91.66% <100.00%> (+11.58%)`	⬆️
pymc/distributions/logprob.py	`100.00% <100.00%> (ø)`
... and 38 more

ricardoV94 · 2022-11-15T11:16:55Z

pymc/data.py

 from aesara.tensor.type import TensorType
 from aesara.tensor.var import TensorConstant, TensorVariable

 import pymc as pm

-from pymc.aesaraf import convert_observed_data
+from pymc.aesaraf import at_rng, convert_observed_data


We can also get rid of at_rng, as the seeds created by it will be overwritten in compile_pymc anyway. Just import RandomStream wherever it's needed.

pymc/model.py

pymc/tests/test_data.py

ricardoV94 · 2022-11-15T18:36:18Z

I think I saw that kernel crash in my PR, perhaps an out of bounds indexing?

ferrine · 2022-11-16T07:28:58Z

no idea why this windows test fails

ferrine · 2022-11-16T07:31:10Z

hmm, that could be the random stream!

ferrine · 2022-11-17T09:12:16Z

The issue that blocks the PR seems to be in aesara or I'm doing something wrong aesara-devs/aesara#1301

ferrine · 2022-11-17T10:20:13Z

Maybe @brandonwillard or @ricardoV94 have any insights?

michaelosthege · 2022-11-19T17:16:37Z

@ferrine are some of these changes independent and could be merged via a separate PR already?

ricardoV94 · 2022-11-19T17:30:10Z

Not really... The change led to some scan issues that need to be addressed otherwise it would lead to a regression

ferrine · 2022-11-21T12:12:30Z

Marked this as draft util we figure out the scan issue

pymc/tests/variational/test_inference.py

pymc/util.py

pymc/variational/approximations.py

pymc/distributions/logprob.py

pymc/variational/opvi.py

ricardoV94 · 2022-11-24T10:14:49Z

pymc/aesaraf.py

 from aesara.scalar.basic import Cast
 from aesara.tensor.basic import _as_tensor_variable
 from aesara.tensor.elemwise import Elemwise
+from aesara.tensor.random import RandomStream


Can we remove the whole at_rng and set_rng related utilities below from this file?

https://github.com/pymc-devs/pymc/blob/12f4cd73bae95c42a28c2f07f6aab60071fd2ca1/pymc/aesaraf.py#L872-L912

set_at_rng is called by pm.sampling.parallel._Process._start_loop.

In the interest of the v5 release I think we should take care of that later

We can, but it's only used by VI

ricardoV94 · 2022-12-07T09:09:41Z

@ferrine What do you think of doing the Minibatch API breaking change for V5, and fixing the MRG later, as that is not user facing anyway.

ferrine · 2022-12-09T06:43:49Z

@ferrine What do you think of doing the Minibatch API breaking change for V5, and fixing the MRG later, as that is not user facing anyway.

That makes to me, we better get back to it later, than get stuck right now. It is also reasonable to improve minibatch API and VI experience later and reduce to something working now

pymc/data.py

pymc/tests/test_data.py

michaelosthege · 2022-12-11T11:51:18Z

I marked the VI test_profile tests as XFAIL and fixed flakyness of the Moyal logcdf test which is also currently ❌ on main.

michaelosthege · 2022-12-11T14:21:40Z

@ferrine I squashed the commits so we can rebase-merge this.

It looks like you have not set up your @pymc-devs.org email address in your GitHub account. That's why your pic doesn't shown on the commits..

pymc/tests/distributions/test_continuous.py

michaelosthege

🥳🚀

ricardoV94 · 2022-12-12T07:20:17Z

@ferrine can you have one last look and merge it if you agree with the changes?

Stale

ferrine · 2022-12-12T09:43:13Z

I had one more look at the PR changes and I don't have anything to add to it. I'll merge once tests pass.

Closes #4523 Closes #6277

michaelosthege · 2022-12-12T13:05:26Z

Just fixuped the edit of the comment into the moyal logcdf commit.

Didn't want to waste CI, so I merged it right away =)

ricardoV94 reviewed Nov 15, 2022

View reviewed changes

pymc/variational/opvi.py Show resolved Hide resolved

ricardoV94 added VI Variational Inference enhancements v4 major Include in major changes release notes section labels Nov 15, 2022

ricardoV94 reviewed Nov 15, 2022

View reviewed changes

pymc/model.py Outdated Show resolved Hide resolved

ricardoV94 changed the title ~~Fix rng streams~~ Refactor MiniBatch and stop using deprecated MRG sampler Nov 15, 2022

ricardoV94 mentioned this pull request Nov 15, 2022

Use default Aesara RandomStream for VI #6273

Closed

ricardoV94 reviewed Nov 15, 2022

View reviewed changes

pymc/tests/test_data.py Outdated Show resolved Hide resolved

ferrine force-pushed the fix-rng-streams branch from b16d9c4 to b66929f Compare November 16, 2022 06:02

ferrine added this to the v4.4.0 milestone Nov 16, 2022

michaelosthege modified the milestones: v4.4.0, v5.0.0 Nov 19, 2022

ferrine changed the title ~~Refactor MiniBatch and stop using deprecated MRG sampler~~ WIP Refactor MiniBatch and stop using deprecated MRG sampler Nov 21, 2022

ferrine marked this pull request as draft November 21, 2022 12:12

ricardoV94 reviewed Nov 22, 2022

View reviewed changes

pymc/tests/variational/test_inference.py Outdated Show resolved Hide resolved

ricardoV94 previously requested changes Nov 24, 2022

View reviewed changes

ricardoV94 changed the title ~~WIP Refactor MiniBatch and stop using deprecated MRG sampler~~ Refactor MiniBatch and stop using deprecated MRG sampler Nov 29, 2022

ferrine force-pushed the fix-rng-streams branch from 1e3b8da to 3acef49 Compare December 9, 2022 13:30

ferrine marked this pull request as ready for review December 9, 2022 16:53

michaelosthege reviewed Dec 9, 2022

View reviewed changes

ferrine force-pushed the fix-rng-streams branch 2 times, most recently from a3093be to 9a19ef4 Compare December 10, 2022 18:50

michaelosthege force-pushed the fix-rng-streams branch 2 times, most recently from 75b697b to f3dde69 Compare December 11, 2022 14:20

michaelosthege force-pushed the fix-rng-streams branch from f3dde69 to 5c86134 Compare December 11, 2022 14:22

michaelosthege requested a review from ricardoV94 December 11, 2022 14:26

michaelosthege mentioned this pull request Dec 11, 2022

BUG: Flaky test test_truncation_specialized_op #6384

Closed

michaelosthege force-pushed the fix-rng-streams branch from 5c86134 to 781fdef Compare December 11, 2022 14:42

ricardoV94 reviewed Dec 11, 2022

View reviewed changes

pymc/tests/distributions/test_continuous.py Outdated Show resolved Hide resolved

michaelosthege mentioned this pull request Dec 11, 2022

BUG: Flaky test test_truncation_continuous_random #6385

Closed

michaelosthege approved these changes Dec 11, 2022

View reviewed changes

michaelosthege and others added 2 commits December 12, 2022 14:01

Fix flaky Moyal logcdf test

be5722c

Refactor Minibatch and stop using MRG_RandomStream in VI

82dfbf6

Closes #4523 Closes #6277

michaelosthege force-pushed the fix-rng-streams branch from 1218d0c to 82dfbf6 Compare December 12, 2022 13:02

michaelosthege merged commit adf6fff into main Dec 12, 2022

michaelosthege deleted the fix-rng-streams branch December 12, 2022 13:04

ricardoV94 mentioned this pull request Dec 14, 2022

Remove global RandomStream #6396

Merged

Uh oh!

Refactor MiniBatch and stop using deprecated MRG sampler #6304

Refactor MiniBatch and stop using deprecated MRG sampler #6304

Uh oh!

Conversation

ferrine commented Nov 15, 2022 • edited by ricardoV94 Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ricardoV94 commented Nov 15, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ricardoV94 commented Nov 15, 2022

Uh oh!

Uh oh!

codecov bot commented Nov 15, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

ricardoV94 Nov 15, 2022

Choose a reason for hiding this comment

Uh oh!

ferrine Nov 15, 2022

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

ricardoV94 commented Nov 15, 2022

Uh oh!

ferrine commented Nov 16, 2022

Uh oh!

ferrine commented Nov 16, 2022

Uh oh!

ferrine commented Nov 17, 2022

Uh oh!

ferrine commented Nov 17, 2022

Uh oh!

michaelosthege commented Nov 19, 2022

Uh oh!

ricardoV94 commented Nov 19, 2022

Uh oh!

ferrine commented Nov 21, 2022

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ricardoV94 Nov 24, 2022

Choose a reason for hiding this comment

Uh oh!

michaelosthege Dec 11, 2022

Choose a reason for hiding this comment

Uh oh!

ricardoV94 Dec 12, 2022

Choose a reason for hiding this comment

Uh oh!

ricardoV94 commented Dec 7, 2022

Uh oh!

ferrine commented Dec 9, 2022

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

michaelosthege commented Dec 11, 2022

Uh oh!

michaelosthege commented Dec 11, 2022

Uh oh!

Uh oh!

michaelosthege left a comment

Choose a reason for hiding this comment

Uh oh!

ricardoV94 commented Dec 12, 2022

Uh oh!

ferrine commented Dec 12, 2022

Uh oh!

michaelosthege commented Dec 12, 2022

ferrine commented Nov 15, 2022 •

edited by ricardoV94

Loading

ricardoV94 commented Nov 15, 2022 •

edited

Loading

codecov bot commented Nov 15, 2022 •

edited

Loading