Numpyro: vectorized elbo does not work with beta bernoulli model

[fehiepsi]

As observed in #413, when running SVI in beta bernoulli with num_particles=2, we get the error

NotImplementedError: Forward-mode differentiation rule for 'while' not implemented

This seems like a JAX issue but we need to isolate the bug to report upstream.

[TuanNguyen27]

@fehiepsi from a quick glance the error traces to here: https://github.com/google/jax/blob/master/jax/interpreters/ad.py#L217

[fehiepsi]

Thanks, @TuanNguyen27 ! I believe that this error comes from while_loop of gamma sampler in JAX, but it is a bit hard to make a small repro script. Anyway, this issue should not block your PR, feel free to complete it and mark failing tests as xfail.

[fehiepsi]

Just want to put a context: this issue is tracked at https://github.com/google/jax/issues/1789.

[fehiepsi]

This is fixed in the latest jax release.