Fix issues in pathfinder wrapper

[ricardoV94]

Closes #97
Closes #250

We were doing senseless things, like drawing from an isometric MvNormal with mean two as the initial point of the VI fit and then discarding such fit? This killed my poor RAM :)

Also I am not sure what the inference_loop was doing. I am not familiar with the internals of blackjax (and the documentation references a version that hasn't yet been released), but I think we just want to do init -> sample? Can someone confirm this (@junpenglao @zaxtax?)

Pathfinder doesn't seem to work particularly well with the 8 school example. It's very sensitive to the initial value of mu, and the default converges to something that looks wrong, see gist. Because of this I don't know what better thing to assert in the test.

Also the pymc-examples is broken (it is not using mu nor tau :P): https://www.pymc.io/projects/examples/en/latest/variational_inference/pathfinder.html

[twiecki]

Ouch, well these all sounds like fixes to me.

[junpenglao]

I think it is better to depends on blackjax-nightly (which would soon be release as blackjax v1), and use

pathfinder = blackjax.kernels.pathfinder(logdensity_fn)
state, _ = pathfinder.approximate(rng_key, ...)
_, rng_key = random.split(rng_key)
samples, _ = pathfinder.sample(rng_key, state, 5_000)

per https://blackjax-devs.github.io/sampling-book/algorithms/pathfinder.html

[ricardoV94]

I think it is better to depends on blackjax-nightly (which would soon be release as blackjax v1), and use

How soon is that soon? Would it be too much of a PITA to get another pre v1 release if that's faster?

[junpenglao]

We just release Blackjax v1

[twiecki]

Can we merge?

[ricardoV94]

I would like to have a model where this actually works as a test case, but since we are in experimental maybe that's fine.

[ricardoV94]

Needs to be rebased for tests to pass