Any interest in a conda-forge recipe?

[jiawen]

I've gotten pretty spoiled by conda install <package>. Is there any interest in adding a recipe on conda-forge for flax? There is already one for jax and jaxlib.

I've never written a conda-forge recipe before but it seems straightforward. I'm not sure I want to sign up to be maintainer though :).

[avital]

Can you share more? What does conda add on top of pip install? I've used conda in the past for packages with binary dependencies, but flax doesn't have any.

[jheek]

I can see why it's nice to stick to a single package manager. We could probably make this part of our GitHub release action so it doesn't require any extra steps during a release

[jiawen]

Yeah, at this point, I try to use conda-forge basically everything. On a new machine, I just do conda install [big list of packages] and am done with it. In my experience, its dependency tracking seems to be more robust than pip - although it's true that I've only had issues when the Python package I'm trying to install had native dependencies that it somehow can't find.

@jheek that would be fantastic!

[BertrandRdp]

Hi @jiawen do you want to give it a shot by modifying what needs to be in actions in our .github directory? let us know if you need something specific.

[jheek]

We can't really accept contributions for setting this up I think. Because you need to setup some secrets to authenticate with the conda package repository

[rwightman]

FYI, the jaxlib conda-forge is CPU only and that's unlikely to change nearterm (jax-ml/jax#189). I'm not sure if that consideration factored into the desire for a Flax conda-forge recipe.

One of the big benefits of a conda based release is the better support for including a full set of binary deps. PyTorch is a great example of this as their conda releases include all cuda/cudnn lib deps so you can install on different platforms without worrying about your system cuda/cudnn versions, only need to make sure your GPU driver is recent enough.

It'd be awesome to have an option to install a full jax + jaxlib (GPU, packaged with CUDA/CUDNN) + flax via Conda but it is a significant maintenance undertaking and would require JAX team to take on the hardest part (jaxlib). It's unlikely the conda-forge team will have the resources to help there.

For now I just use docker as juggling JAX related pypi packages, w/ TF + Tensorboard with working cuda/cudnn/nccl versions has been problematic.