BF: fix ensemble mcmc run after warmup

[amifalk]

Fixes #1916. The ensemble kernel relies on the batched shape of the rng_key to infer the number of chains but subsets the batched key into a singleton internally. We replicate this behavior in the mcmc.run method when _warmup_state exists by injecting a singleton key into the ensemble state.

[amifalk]

@fehiepsi Looks like linting is failing on files unrelated to this PR. Do you want me to fix them here?

[OlaRonning]

It looks like they are fixed in #1917, which is merged. If you rebase, the lint should pass.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[fehiepsi]

Could you rebase the change? You can use something like

git checkout master
git fetch origin master
git rebase origin master
git checkout ensemble-warmup
git merge master
git push [remote to your fork] ensemble-warmup

[fehiepsi]

Thanks @amifalk!