jax + jaxlib in pyproject.toml -- or just jax?

[femtomc]

Hi all!

I'm attempting to figure out the right way to expose my library to users, and how I should handle dependencies.

Right now, I just specify jax as a dependency (with version bounds) -- and I was under the assumption that different variants of jaxlib support different things (e.g. different backends). Is this wrong?

More generally: how should I handle jax as a dependency? Should I keep what I've done so far, or should I also specify jaxlib?

[jakevdp]

jaxlib is required to run any jax code, but jax does not pin jaxlib in its requirements because different jaxlib libraries are needed for different hardware (cpu, gpu, tpu, etc.) and Python's dependency tracking system does not have the granularity of expression for that. JAX itself addressed that via the installation instructions.

If you have a jax-dependent library which is meant to be run on multiple platforms, you inherit these difficulties. If you list jaxlib as a hard dependency in your pyproject.toml, it may prevent users from installing the correct jaxlib version for their system. So I'd recommend following JAX, and not listing jaxlib in pyproject.toml, but rather be clear in your package installation instructions that jaxlib must be installed separately.

You might also think about defining extra installation commands, e.g. yourproject[cpu] that installs the same jaxlib as jax[cpu]. It's not a super-clean solution, but given the limitations of Python's dependency specification, I think it's the best you can do. What do you think?

[femtomc]

Thanks Jake! Seems like I have the right idea with what I’m currently doing, sounds good.