Conflicting dependencies for jax[cuda11-pip]==0.4.14

[gbacon]

On commit f926148, running conda env create -f scripts\gpu_environment.yml fails with

INFO: pip is looking at multiple versions of jax[cuda11-pip] to determine which version is compatible with other requirements. This could take a while.

The conflict is caused by:
    optax 0.1.7 depends on jaxlib>=0.1.37
    distrax 0.1.4 depends on jaxlib>=0.1.67
    chex 0.1.82 depends on jaxlib>=0.1.37
    jax[cuda11-pip] 0.4.14 depends on jaxlib==0.4.14+cuda11.cudnn86; extra == "cuda11_pip"

To fix this you could try to:
1. loosen the range of package versions you've specified
2. remove package versions to allow pip attempt to solve the dependency conflict

Pip subprocess error:
WARNING: jax 0.4.14 does not provide the extra 'cuda11-pip'
ERROR: Cannot install -r ...\EasyLM\scripts\condaenv.ur0dg_4_.requirements.txt (line 5), -r ...\EasyLM\scripts\condaenv.ur0dg_4_.requirements.txt (line 6), -r ...\EasyLM\scripts\condaenv.ur0dg_4_.requirements.txt (line 7) and jax[cuda11-pip]==0.4.14 because these package versions have conflicting dependencies.
ERROR: ResolutionImpossible: for help visit https://pip.pypa.io/en/latest/topics/dependency-resolution/#dealing-with-dependency-conflicts

failed

CondaEnvException: Pip failed

This is using python Python 3.11.5, Anaconda3-2023.09, and conda 23.7.4.

After removing the 3.10 dependency for Python on line 5 of scripts/gpu_environment.yml, running conda env remove -n EasyLM, and rerunning the aforementioned create command fails with

Pip subprocess output:
Looking in indexes: https://pypi.org/simple, https://download.pytorch.org/whl/cu118
Looking in links: https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
Collecting jax==0.4.14 (from jax[cuda11_pip]==0.4.14->-r ...\EasyLM\scripts\condaenv.wv5z019b.requirements.txt (line 3))
  Using cached jax-0.4.14.tar.gz (1.3 MB)
  Installing build dependencies: started
  Installing build dependencies: finished with status 'done'
  Getting requirements to build wheel: started
  Getting requirements to build wheel: finished with status 'done'
  Preparing metadata (pyproject.toml): started
  Preparing metadata (pyproject.toml): finished with status 'done'
Collecting flax==0.7.0 (from -r ...\EasyLM\scripts\condaenv.wv5z019b.requirements.txt (line 4))
  Using cached flax-0.7.0-py3-none-any.whl.metadata (9.9 kB)
Collecting optax==0.1.7 (from -r ...\EasyLM\scripts\condaenv.wv5z019b.requirements.txt (line 5))
  Using cached optax-0.1.7-py3-none-any.whl.metadata (13 kB)
Collecting distrax==0.1.4 (from -r ...\EasyLM\scripts\condaenv.wv5z019b.requirements.txt (line 6))
  Using cached distrax-0.1.4-py3-none-any.whl.metadata (13 kB)
Collecting chex==0.1.82 (from -r ...\EasyLM\scripts\condaenv.wv5z019b.requirements.txt (line 7))
  Using cached chex-0.1.82-py3-none-any.whl.metadata (17 kB)
Collecting transformers==4.31.0 (from -r ...\EasyLM\scripts\condaenv.wv5z019b.requirements.txt (line 8))
  Using cached transformers-4.31.0-py3-none-any.whl.metadata (116 kB)

Pip subprocess error:
ERROR: Could not find a version that satisfies the requirement torch==2.0.1 (from versions: none)
ERROR: No matching distribution found for torch==2.0.1

failed

CondaEnvException: Pip failed