Waste Free SMC available for adaptive tempered and tempered SMC.

[ciguaran]

Although SMC base.py was ready for waste-free, tempered and adaptive tempered are coupled to using classic SMC (iterate MCMC and keep the last sample). This PR makes Waste-free available at all levels of the hierarchy.
#710

[codecov]

Codecov Report

Attention: Patch coverage is 94.87179% with 2 lines in your changes missing coverage. Please review.

Project coverage is 98.67%. Comparing base (20666de) to head (957b958).
Report is 16 commits behind head on main.

Files	Patch %	Lines
blackjax/smc/waste_free.py	92.59%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #721      +/-   ##
==========================================
- Coverage   98.72%   98.67%   -0.05%     
==========================================
  Files          59       60       +1     
  Lines        2746     2803      +57     
==========================================
+ Hits         2711     2766      +55     
- Misses         35       37       +2     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[AdrienCorenflos]

By the way, whilst we are on that, I remember being a bit dubious about waste free for parallel environments: when you have a lot of processors the analysis in the paper does not hold anymore.
@hai-dang-dau have you ever done some back of the envelope calculations on this? I'd like to be able to recommend people to switch off waste free (and maybe increase the number of intermediary distributions in exchange?) if they have a certain number of processors.

[hai-dang-dau]

It's true that waste-free SMC shines best when the number of processors is small relative to the number of particles. (Say, a dozen of processors for several thousand of particles.) I think this is the case for most practical applications.

If your number of particles is comparable to that of processors (say, 500 processors and 500 particles), then yes standard SMC is better.

However, if you have that many processors, you probably have a lot of memory as well. When you increase the number of particles to about a million (still with 500 CPUs), waste-free starts to be preferable again.

[AdrienCorenflos]

Well, not really: a commercial grade GPU has fairly little memory so it's not all that black and white. I'll have a think.

…

On Fri, 16 Aug 2024, 11:09 hai-dang-dau, ***@***.***> wrote: It's true that waste-free SMC shines best when the number of processors is small relative to the number of particles. (Say, a dozen of processors for several thousand of particles.) I think this is the case for most practical applications. If your number of particles is comparable to that of processors (say, 500 processors and 500 particles), then yes standard SMC is better. However, if you have that many processors, you probably have a lot of memory as well. When you increase the number of particles to about a million (still with 500 CPUs), waste-free starts to be preferable again. — Reply to this email directly, view it on GitHub <#721 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AEYGFZ5MOG624WOTBJF2BQTZRXFWRAVCNFSM6AAAAABMSYWGTWVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDEOJTGIZDQOBSHA> . You are receiving this because you commented.Message ID: ***@***.***>

[hai-dang-dau]

In my previous post, by "processors" I meant the number of MCMC chains that can be run on parallel. Although GPU has a lot of parallel processing capacity, I'm not sure how this translates into the number of MCMC chains. The delicate point is that one MCMC chain itself involves a lot of parallelizable operations (ie matrix multiplications for calculating the likelihood or for generating preconditioned proposals). It is these primitive operations that are usually parallelized automatically by JAX/XLA, not the higher level operations that were taken into account in my reasoning above.

Having said that, it's true that we developed and tested waste-free SMC with CPUs in mind; so it'd definitely be relevant to have recommendations tailored to GPU-accelerated environment.

[junpenglao]

Minor nit

[junpenglao]

Feel free to merge after fixing the code style

[ciguaran]

@junpenglao the review gets dismissed whenever we commit something new, even if it is just github's commiting a reviewer suggestion

[junpenglao]

Thanks!