Proper way to expose GPU-enabled JAX via a library which relies on JAX

[femtomc]

Hi all!

I'm developing a library which I hope to share with users - I'm using poetry for dependency management within my pyproject.toml.

What's the best practices surrounding facilitating users to use a GPU-enabled variant of JAX when they install such a library? Has anyone investigated this before?

I think roughly a user can always just overload their current Python environment by first installing my library, and then using pip to install a GPU-enabled variant of JAX.

I'm curious if that's basically how people solve this issue -- and if there's nothing better.