Installing JAX across MacOS and Linux.

The markers to specify separate versions/sources for different platforms are not working as I would've expected.

jaxlib

With jaxlib, I was able to resolve the problem via a hacky workaround.

The simple version did not work:

jaxlib = [
    {version = "^0.3.22", platform = "darwin"},
    {version = "^0.3.22+cuda11.cudnn82", platform = "linux"}
]

With this dependency spec, poetry v1.2.2 will install the +cuda version of the package regardless of the platform.

This seems to be an example of the issue discussed in https://github.com/python-poetry/poetry/issues/6710 and fixed in https://github.com/python-poetry/poetry-core/pull/497

I worked around the issue by specifying like:

jaxlib = [
    {version = "0.3.20", platform = "darwin", source="pypi"},
    {version = "^0.3.22+cuda11.cudnn82", platform = "linux"}
]

Since the version numbers are now different, the darwin install works correctly. The source="pypi" was also necessary.

jax

With jax, I'm unable to get anything to work in a cross-platform way. As one particular example, I pinned versions to try to get something like the jaxlib hack to work:

jax = [
    {extras = ["cpu"], version = "0.3.21", platform="darwin", source="pypi"},
    {extras = ["cuda"], version = "0.3.22", platform="linux"}
]

In this case, the resolver seems to decide that I want to depend on both version:

Because jax (0.3.22) depends on jaxlib (0.3.22+cuda11.cudnn82)
 and jax (0.3.21) depends on jaxlib (0.3.20), jax (0.3.22) is incompatible with jax (0.3.21).
So, because confirm depends on both jax (0.3.21) and jax (0.3.22), version solving failed.

details

I previously posted: https://github.com/python-poetry/poetry/issues/6955 and in the comment there, I thought I had fixed this issue, but that was just from looking at the poetry.lock file. Actual testing demonstrated I had been wrong.

Versions:

• Poetry version: 1.2.2
• Python version: 3.10.6
• OS version and name: Ubuntu 22.04
• pyproject.toml:

[tool.isort]
profile = "black"

[tool.poetry]
name = "confirm"
version = "0.1.0"
description = ""
authors = ["Your Name <you@example.com>"]

[tool.poetry.dependencies]
python = "~3.10"
numpy = "^1.23.4"
scipy = "^1.9.3"
typer = "^0.6.1"
jax = [
    {extras = ["cpu"], version = "0.3.21", platform="darwin", source="pypi"},
    {extras = ["cuda"], version = "0.3.22", platform="linux"}
]

[tool.poetry.group.test.dependencies]
pytest = "^7.2.0"
pre-commit = "^2.20.0"

[[tool.poetry.source]]
name = "jax"
url = "https://storage.googleapis.com/jax-releases/jax_cuda_releases.html"
default = false
secondary = false

[build-system]
requires = ["poetry-core>=1.0.0"]
build-backend = "poetry.core.masonry.api"

Start with the provided pyproject.toml and reproduce dependency resolution using the reported Poetry 1.2.2, Python 3.10.6, and platform-specific jax and jaxlib sources. Trace how platform markers, versions, extras, and sources are combined; done means compatible dependencies resolve separately on macOS and Linux without requiring mismatched versions.