Avoid recompiling NumPyro models where possible

[null-a]

If we do something like:

model = define_model(...)
fit1 = model.run_algo('nuts', ...)
fit2 = model.run_algo('nuts', ...)

... then each call to run_algo will invoke NumPyro's model compilation. It ought to be possible to avoid this (in cases where the data shape doesn't change between calls) by caching the MCMC instance and reusing it on subsequent calls.

[neerajprad]

(in cases where the data shape doesn't change between calls)

Even if the data shape changes, but we can commit to a certain max size, it might be possible to use a masked distribution to enable this. We haven't needed a masked distribution in numpyro yet, but this may be a good use case. cc. @fehiepsi

[null-a]

The feature is available in NumPyro >=0.2.2.