Investigate within-chain parallelization in numpyro

[dylanhmorris]

Goal

Understand what numpyro support exists for within-chain parallelization.

Details

• The plate context can be used to mark conditionally independent sampling operations.
• Are there ways then to ask the compiler to parallelize things within a single MCMC chain, as is possible in Stan?

See also

Related discussion in EpiAware CDCgov/Rt-without-renewal#254

[dylanhmorris]

@seabbs reports that it looks possible with JAX but not implemented in numpyro's bundled samplers. So if we want/need this future the solution may be to do sampling via blackjax or bayeux (see also #361).

In general, perhaps worth considering whether to separate PPL choice and sampler choice.

[seabbs]

*but that he does not think it is possible based on current evidence (I am heavily caveating as I really think it must be I just can't find any evidence).

[seabbs]

Okay some actual evidence: jax-ml/jax#1408

I think numpyro is using pmap to run multiple chains but that doesn't mean it couldn't be used within one chain (unless you can't stack them which I would hope you can but probably can't). That thread also shows a use of vmap in blackjax but its not clear it 1. would work in numpyro and 2. can be stacked with pmap.

So I am still leaning to no it can't right now but I think above gives some room for exploration that I didn't have a handle on before.

[seabbs]

The numpyro docs for setting the XLA options for this say they don't understand what it will do so that is encouraging.

[seabbs]

https://num.pyro.ai/en/latest/utilities.html#numpyro.util.set_host_devices

[damonbayer]

https://num.pyro.ai/en/latest/utilities.html#numpyro.util.set_host_devices

We already use this to enable between-chain parallelization.

[seabbs]

Its the warning that I was drawing attention to. The key question is it hard coded to only allow across chain parallelization or can you stack.