Numpyro: installation / deployment to conda-forge

[geoHeil]

Is anyone working on making this library available via conda-forge?

[fehiepsi]

Hi @geoHeil, we don't have a plan for making NumPyro available in conda-forge. NumPyro depends on JAX, which does not have a conda build. In addition, we don't have the bandwidth to maintain conda builds for Pyro packages. NumPyro is purely Python and only has a few dependencies, so using pip is enough I guess.

[geoHeil]

Are you certain?

• https://github.com/conda-forge/jaxlib-feedstock
• https://github.com/conda-forge/jax-feedstock

both are registered as feedstocks.

As far as I understand it, for now, this is only for the CPU versions (yet) on conda, but I might be wrong.

[neerajprad]

We had some discussion around this in Pyro - https://github.com/pyro-ppl/pyro/issues/1494, which may be relevant. Since both Pyro and NumPyro are just python libraries, you should be able to use pip in your conda environment or install it through an environment.yml file without issues. Are there any specific reasons to have a separate conda package?

[geoHeil]

Indeed:

channels:
  - conda-forge
  - defaults
dependencies:
  - jax
  - jaxlib
  - pip:
      - git+https://github.com/pyro-ppl/[email protected]

something along these lines can be set. But from the logs it looks like pip then installs a couple of libraries like jax again.

[fehiepsi]

both are registered as feedstocks.

Interesting, thanks! I didn't know that some contributors are maintaining jax/jaxlib builds in conda-forge. This JAX thread discusses the progress so far. It seems like currently, only CPU version is available.

pip then installs a couple of libraries like jax

I think this issue happens because numpyro releases depend on specific JAX versions. We might consider removing it when JAX 0.2 is released. Until then, jax APIs might change between each release and will break NumPyro users' code.

I think you can convert pip install https://github.com/pyro-ppl/pyro-api/archive/master.zip to conda version to install the master branch of NumPyro, which is tested for the latest JAX release. Maybe git+https://github.com/pyro-ppl/numpyro.git@master?

[geoHeil]

How stable is the numpyro master? I thought it is better to use the latest 0.2.4 release.

[fehiepsi]

We tried to keep up with the latest JAX release, but there is a serious bug in JAX that is likely to be fixed in the next JAX release: we can't reshape samples from multi-chain MCMC. The remaining functionalities are pretty stable IMO.

[neerajprad]

I suppose you can also pin the jax and jaxlib versions in your environment.yml file to the same as in numpyro's setup.py. The reason why we have needed to pin to a jax/jaxlib release in the past is that we found small changes to jax's api that would end up breaking compatibility, though as things become more stable, we would be able to pin to a min version instead.

[fehiepsi]

I marked this issue as "help wanted" in case some NumPyro users want to create and maintain a conda-forge build for NumPyro. Because some versions of jax/jaxlib are available in conda-forge, I guess registering NumPyro there might not be so complicated.

[freddyaboulton]

@fehiepsi I have a PR to the sample-recipes repo to add numpyro to conda. Let me know if you want me to make any changes (add maintainers) before opening it for review.