Add support for torch 2.11 and jax 0.9.1/0.9.2, remove support for jax 0.6.0/0.6.1 (#213)

Author: frostedoyster

.github/workflows/build-jax-wheels.yml

@@ -51,7 +51,7 @@ jobs:
     runs-on: ${{ matrix.os }}
     strategy:
       matrix:
-        jax-version: ["0.6.0", "0.6.1", "0.6.2", "0.7.0", "0.7.1", "0.7.2", "0.8.0", "0.8.1", "0.8.2", "0.8.3", "0.9.0"]
+        jax-version: ["0.6.2", "0.7.0", "0.7.1", "0.7.2", "0.8.0", "0.8.1", "0.8.2", "0.8.3", "0.9.0", "0.9.1", "0.9.2"]
         os: [ubuntu-22.04, macos-14, ubuntu-22.04-arm]
         include:
           - name: x86_64 Linux

.github/workflows/build-torch-wheels.yml

@@ -51,7 +51,7 @@ jobs:
     runs-on: ${{ matrix.os }}
     strategy:
       matrix:
-        pytorch-version: ["2.3", "2.4", "2.5", "2.6", "2.7", "2.8", "2.9", "2.10"]
+        pytorch-version: ["2.3", "2.4", "2.5", "2.6", "2.7", "2.8", "2.9", "2.10", "2.11"]
         os: [ubuntu-22.04, macos-14, ubuntu-22.04-arm]
         include:
           - os: ubuntu-22.04

docs/requirements.txt

@@ -8,7 +8,7 @@ torch
 
 # jax[cpu], because python -m pip install jax, which would be triggered
 # by the main package's dependencies, does not install jaxlib
-jax[cpu] <0.6
+jax[cpu]
 
 # metatensor and metatensor-torch for the metatensor API
 metatensor-torch

tox.ini

@@ -96,7 +96,7 @@ commands =
 
     # Install this one manually. Listing it in the deps list above does not install jaxlib.
     # Note: jax[cuda12] is not available on Windows and MacOS.
-    bash -c 'command -v nvcc && python -m pip install -U "jax[cuda12]<0.7.2" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html || python -m pip install -U "jax[cpu]<0.7.2"'
+    bash -c 'command -v nvcc && python -m pip install -U "jax[cuda12]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html || python -m pip install -U "jax[cpu]"'
 
     pip install {[testenv]pip_install_flags} .
     pytest python {posargs}
@@ -112,7 +112,7 @@ deps =
     absl-py # jax uses this package but not installed by jax[cpu]
     pytest
     equinox
-    jax[cpu]==0.6.0
+    jax[cpu]==0.6.2
 
 allowlist_externals =
     bash
@@ -140,7 +140,7 @@ deps =
 commands =
     # Install this one manually. Listing it in the deps list above does not install jaxlib.
     # Note: jax[cuda12] is not available on Windows and MacOS.
-    bash -c 'command -v nvcc && python -m pip install -U "jax[cuda12]<0.7.2" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html || python -m pip install -U "jax[cpu]<0.7.2"'
+    bash -c 'command -v nvcc && python -m pip install -U "jax[cuda12]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html || python -m pip install -U "jax[cpu]"'
 
     pip install {[testenv]pip_install_flags} .
     pip install {[testenv]pip_install_flags} ./sphericart-torch