ContextualVersionConflict when using from_numpyro

Describe the bug This is the strangest thing and I haven’t been able to entirely work out what’s going on but perhaps y’all have thoughts!

When using Google colab, I install numpyro with --no-deps so that it doesn’t try to update JAX to the CPU-only version. Then, everything runs fine (the sampler runs and I get my trace) until I execute az.from_numpyro(), at which point it blows up with the following exception:

ContextualVersionConflict                 Traceback (most recent call last)

<ipython-input-3-0b6cc67427eb> in <module>()
----> 1 data = az.from_numpyro(mcmc)

9 frames

/usr/local/lib/python3.6/dist-packages/pkg_resources/__init__.py in resolve(self, requirements, env, installer, replace_conflicting, extras)
    773                 # Oops, the "best" so far conflicts with a dependency
    774                 dependent_req = required_by[req]
--> 775                 raise VersionConflict(dist, req).with_context(dependent_req)
    776 
    777             # push the new requirements onto the stack

ContextualVersionConflict: (jax 0.2.6 (/usr/local/lib/python3.6/dist-packages), Requirement.parse('jax==0.2.3'), {'numpyro'})

This doesn’t really seem to be an ArviZ issue, but everything is fine until I try to use ArviZ so perhaps you have thoughts.

To Reproduce In a factor reset colab environment I run the following and get the above exception on the last line:

!pip install --no-deps numpyro
!pip install arviz

import arviz as az
import numpyro
import jax

def model():
    numpyro.sample("x", numpyro.distributions.Normal(0, 1))

mcmc = numpyro.infer.MCMC(numpyro.infer.NUTS(model), 1000, 1000, num_chains=4)
mcmc.run(jax.random.PRNGKey(34923))

data = az.from_numpyro(mcmc)

Expected behavior I would expect this exception to be thrown earlier if it is actually a problem, or for ArviZ to run as expected!

[ColCarroll]

Wow! Can repro, and can confirm that the same behavior happens with PRNGKey 86074.

More seriously, the bug happens here, when trying to grab the library version to stamp the inference data. This definitely shouldn’t raise, and indicates that we’re doing something wrong with pkg_resources, and maybe with lazy-loading libraries.

Shorter term, if you import numpyro before arviz, it works fine.

[OriolAbril]

Is this still happening?