Compare performance of brmp (pyro/numpyro) and brms

[null-a]

[null-a]

Including both CPU and GPU (#14).

[neerajprad]

@null-a : We were running some benchmarks here - https://github.com/pyro-ppl/numpyro/pull/470, and are mostly concluding that right now. From what we can see, it seems very likely that you will observe good performance on all your models. Let us know if there are any particular models that you would like to test out, or have observed some things to be slower than expected. We are also going to be doing a minor release tomorrow which should fix the memory issue, and the issue of model code being recompiled when we run MCMC.

For the latter, recompilation can be avoided by simply calling .run on the same MCMC instance, but the data shape needs to be the same too (XLA will recompile the code if the input data shape changes). I am not sure if the current brmp code is set up to take advantage of avoiding recompilation this way, but if this seems important, we should set up the backend to take advantage of this.

[null-a]

I am not sure if the current brmp code is set up to take advantage of avoiding recompilation this way

No, not yet. (Now #79)

but if this seems important

FWIW, the use case I have involves repeatedly performing inference on a model every time a new data point arrives, so it sounds like this particular case might not benefit from caching.

[fehiepsi]

this particular case might not benefit from caching

This is tricky... I think the proposal from @neerajprad to embed the current data into a holder and provide data_size information to mask out data will work here. The pattern will be something like

data = np.zeros(LIMIT_SIZE)
# assume that currently, data has size `size`
# when a new data point comes
size = size + 1
data[size] = new_data_point
mcmc.run(rng_key, data, size)

and for the model

def model(data, size):
   ...
   current_data = data[:size]  # probably use lax.dynamic_slide_in_dim here
   sample("obs", dist.Normal(mu, sigma), obs=current_data)

WDYT about it? I think this will work for regression cases, but caching won't work if there is a latent variable having shape depend on size...

Probably this will be easier if XLA supports caching with dynamic data size in the future.