Support for extra features like in jax[gpu]

[gl-yziquel]

Hi.

I've been trying to build using python -m build the jax project: https://github.com/google/jax

The build system comes with some extras_require. See here: https://github.com/google/jax/blob/main/setup.py#L88

As of today, I see no way to build wheels for jax[cpu] or jax[gpu], but only for jax without any features, and thus without libjax, which makes the build useless.

I skimmed through the docs and saw nowhere how the equivalent of python -m build --wheel --extras="cpu" could be achieved. Could that extra --extras="cpu" flag be worth implemening ? I believe it should.

[layday]

Extras list runtime dependencies; jaxlib is a runtime dependency of jax. If you need jaxlib, you have to build it separately.

[gl-yziquel]

ok. doing that.