Repository: DPBayes/jax-chacha-prng (language: python)

Issue:
Import breaks with JAX version 0.4.24 due to removed function
Error message from importing `chacha.random`:
```
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/oraisa/Library/Caches/pypoetry/virtualenvs/multiple-syn-datasets-ml-CFH753D6-py3.10/lib/python3.10/site-packages/chacha/random.py", line 28, in <module>
    import chacha.cipher as cc
  File "/Users/oraisa/Library/Caches/pypoetry/virtualenvs/multiple-syn-datasets-ml-CFH753D6-py3.10/lib/python3.10/site-packages/chacha/cipher.py", line 19, in <module>
    from chacha.jax_ops import chacha20_block as _block
  File "/Users/oraisa/Library/Caches/pypoetry/virtualenvs/multiple-syn-datasets-ml-CFH753D6-py3.10/lib/python3.10/site-packages/chacha/jax_ops.py", line 48, in <module>
    xla_client.register_cpu_custom_call_target(
AttributeError: module 'jaxlib.xla_client' has no attribute 'register_cpu_custom_call_target'. Did you mean: 'register_custom_call_target'?
```


Discussion / hints:

Fix patch:
diff --git a/.github/workflows/build_wheels_and_publish.yml b/.github/workflows/build_wheels_and_publish.yml
index 25a1321..8f19981 100644
--- a/.github/workflows/build_wheels_and_publish.yml
+++ b/.github/workflows/build_wheels_and_publish.yml
@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: © 2023 Aalto University
+# SPDX-FileCopyrightText: © 2023 Aalto University, © 2024 Lukas Prediger
 
 name: Build_pypi_wheels
 
@@ -36,12 +36,6 @@ jobs:
       - uses: actions/checkout@v3
         with:
           submodules: 'recursive'
-
-      # - uses: actions/setup-python@v4
-      #   name: Install Python
-      #   with:
-      #     python-version: '3.9'
-
       - name: Set up QEMU
         if: runner.os == 'Linux' && matrix.arch == 'aarch64'
         uses: docker/setup-qemu-action@v2
diff --git a/ChangeLog.txt b/ChangeLog.txt
index f889879..8ce21af 100644
--- a/ChangeLog.txt
+++ b/ChangeLog.txt
@@ -1,3 +1,5 @@
+- 1.4.2:
+    - Hotfix: preventing installation from pulling incompatible jax versions >0.4.14
 - 1.4.1:
     - Fix: native module not built for correct Python verison for MacOS.
     - Fix: Wheels for MacOS no longer build with OpenMP to avoid complicatons with conda environments.
diff --git a/ChangeLog.txt.license b/ChangeLog.txt.license
index 3874c77..339fac9 100644
--- a/ChangeLog.txt.license
+++ b/ChangeLog.txt.license
@@ -1,2 +1,2 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: © 2021 Aalto University
+# SPDX-FileCopyrightText: © 2023 Aalto University, © 2024 Aalto University
diff --git a/README.md b/README.md
index feb0355..79a10c6 100644
--- a/README.md
+++ b/README.md
@@ -94,7 +94,7 @@ JAX version known to be compatible with JAX-ChaCha-PRNG:
 pip install .[compatible-jax]
 ```
 
-JAX-ChaCha-PRNG is currently known to work reliably with JAX versions 0.2.12 - 0.3.25 .
+JAX-ChaCha-PRNG is currently known to work reliably with JAX versions 0.2.12 - 0.4.14 .
 We regularly check the compatible version range, but do not expect new versions of JAX to be immediately tested.
 
 ## Versioning
diff --git a/README.md.license b/README.md.license
index 3874c77..339fac9 100644
--- a/README.md.license
+++ b/README.md.license
@@ -1,2 +1,2 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: © 2021 Aalto University
+# SPDX-FileCopyrightText: © 2023 Aalto University, © 2024 Aalto University
diff --git a/chacha/jax_ops.py b/chacha/jax_ops.py
index 7cbe723..ed8309a 100644
--- a/chacha/jax_ops.py
+++ b/chacha/jax_ops.py
@@ -21,26 +21,31 @@
 
 # importing ShapedArray and dtypes
 try:
-    # pre jax v0.2.14 location
-    from jax.abstract_arrays import ShapedArray  # type: ignore
-    from jax import dtypes  # type: ignore
+    # post jax v0.4.13 (or so) location
+    from jax.core import ShapedArray  # type: ignore
+    from jax._src import dtypes  # type: ignore
 except (AttributeError, ImportError):  # pragma: no cover
-    # post jax v0.2.14 location
     try:
+        # post jax v0.2.14 location
         from jax._src.abstract_arrays import ShapedArray  # type: ignore
         from jax._src import dtypes  # type: ignore
     except (AttributeError, ImportError):  # pragma: no cover
-        raise ImportError("Cannot import ShapedArray and dtypes. "
+        try:
+            # pre jax v0.2.14 location
+            from jax.abstract_arrays import ShapedArray  # type: ignore
+            from jax import dtypes  # type: ignore
+        except (AttributeError, ImportError):  # pragma: no cover
+            raise ImportError("Cannot import ShapedArray and dtypes. "
                           "You are probably using an incompatible version of jax.")
 
 # importing XlaOp
 try:
-    # pre jax v0.2.13 location
-    XlaOp = jax.xla.XlaOp  # type: ignore
-except (AttributeError, ImportError):  # pragma: no cover
     # post jax v0.2.13 location
+    from jaxlib.xla_client import XlaOp  # type: ignore
+except (AttributeError, ImportError):  # pragma: no cover
     try:
-        from jaxlib.xla_client import XlaOp  # type: ignore
+        # pre jax v0.2.13 location
+        XlaOp = jax.xla.XlaOp  # type: ignore
     except (AttributeError, ImportError):  # pragma: no cover
         raise ImportError("Cannot import XlaOp. "
                           "You are probably using an incompatible version of jax.")
diff --git a/chacha/version.py b/chacha/version.py
index d71c874..e7345b8 100644
--- a/chacha/version.py
+++ b/chacha/version.py
@@ -1,9 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: © 2023 Aalto University
+# SPDX-FileCopyrightText: © 2023 Aalto University, © 2024 Lukas Prediger
 
 MAJOR_VERSION = 1
 MINOR_VERSION = 4
-PATCH_VERSION = 1
+PATCH_VERSION = 2
 EXT_VERSION = ""
 POST_VERSION = ""
 
diff --git a/setup.py b/setup.py
index a538ab7..06eaa78 100644
--- a/setup.py
+++ b/setup.py
@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: © 2022 Aalto University
+# SPDX-FileCopyrightText: © 2023 Aalto University, © 2024 Lukas Prediger
 
 import setuptools
 from setuptools import Extension
@@ -84,8 +84,8 @@ def build_extension(self, ext):
 spec.loader.exec_module(version_module)
 
 _jax_version_lower_constraint = ' >= 0.2.12'
-_jax_version_optimistic_upper_constraint = ', <= 2.0.0'
-_jax_version_upper_constraint = ', <= 0.4.8'
+_jax_version_optimistic_upper_constraint = ', <= 0.4.14'
+_jax_version_upper_constraint = ', <= 0.4.14'
 
 _version = version_module.VERSION
 if 'JAX_CHACHA_PRNG_BUILD' in os.environ:
@@ -101,7 +101,7 @@ def build_extension(self, ext):
     long_description_content_type="text/markdown",
     url="https://github.com/DPBayes/jax-chacha-prng",
     packages=setuptools.find_packages(include=['chacha', 'chacha.*']),
-    python_requires='>=3.6',
+    python_requires='>=3.6, <3.12',
     install_requires=[
         "numpy >= 1.16, < 2",
         "deprecation < 3",
@@ -109,7 +109,7 @@ def build_extension(self, ext):
     ],
     extras_require={
         "tests": [
-            f"jax[minimum-jaxlib]",
+            "jax[cpu]",
             "pytest"
         ],
         "compatible-jax": [f"jax{_jax_version_lower_constraint}{_jax_version_upper_constraint}"]


Test patch:
diff --git a/.github/workflows/jax_compatibility_tests.yml b/.github/workflows/jax_compatibility_tests.yml
index 2f76945..734e6e0 100644
--- a/.github/workflows/jax_compatibility_tests.yml
+++ b/.github/workflows/jax_compatibility_tests.yml
@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: © 2023 Aalto University
+# SPDX-FileCopyrightText: © 2023 Aalto University, © 2024 Aalto University
 
 name: Jax Compatibility Tests
 
@@ -33,7 +33,7 @@ jobs:
         python -m pip install pytest
     - name: Install dependencies
       run: |
-        python -m pip install "jax[minimum-jaxlib]==${{ matrix.jax-version }}" "numpy < 1.24"
+        python -m pip install "jax[minimum-jaxlib]==${{ matrix.jax-version }}" "numpy < 1.24" -f https://storage.googleapis.com/jax-releases/jax_releases.html
         python -m pip install .
     - name: Test with pytest
       run: |
@@ -46,8 +46,7 @@ jobs:
       matrix:
         jax-version: [
           0.2.18, 0.2.19, 0.2.20, 0.2.21,
-          0.2.22, 0.2.27, 0.3.1, 0.3.13, 0.3.15,
-          0.3.17, 0.3.23, 0.3.25
+          0.2.22, 0.2.27
         ]
 
     steps:
@@ -64,7 +63,37 @@ jobs:
         python -m pip install pytest
     - name: Install dependencies
       run: |
-        python -m pip install "jax[minimum-jaxlib]==${{ matrix.jax-version }}"
+        python -m pip install "jax[minimum-jaxlib]==${{ matrix.jax-version }}" -f https://storage.googleapis.com/jax-releases/jax_releases.html
+        python -m pip install .
+    - name: Test with pytest
+      run: |
+        pytest tests/
+
+  unittests-with-current-jaxlib:
+    runs-on: ubuntu-20.04
+    strategy:
+      fail-fast: false
+      matrix:
+        jax-version: [
+          0.3.1, 0.3.13, 0.3.15,
+          0.3.17, 0.3.23, 0.3.25, 0.4.14
+        ]
+
+    steps:
+    - uses: actions/checkout@v2
+      with:
+        submodules: true
+    - name: Set up Python
+      uses: actions/setup-python@v2
+      with:
+        python-version: 3.9
+    - name: Install environment
+      run: |
+        python -m pip install --upgrade pip
+        python -m pip install pytest
+    - name: Install dependencies
+      run: |
+        python -m pip install "jax[cpu]==${{ matrix.jax-version }}" -f https://storage.googleapis.com/jax-releases/jax_releases.html
         python -m pip install .
     - name: Test with pytest
       run: |
diff --git a/.github/workflows/python_unittests.yml b/.github/workflows/python_unittests.yml
index 4b47751..0a90eae 100644
--- a/.github/workflows/python_unittests.yml
+++ b/.github/workflows/python_unittests.yml
@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: © 2021 Aalto University
+# SPDX-FileCopyrightText: © 2023 Aalto University, © 2024 Lukas Prediger
 
 name: Python Unittests
 
@@ -37,7 +37,7 @@ jobs:
         flake8 --count --exit-zero --max-complexity=10 --max-line-length=120 --statistics --ignore=E266 chacha tests
     - name: Install dependencies
       run: |
-        python -m pip install -e .[tests]
+        python -m pip install -e .[tests] -f https://storage.googleapis.com/jax-releases/jax_releases.html
     - name: Test with pytest
       run: |
         pytest --cov=chacha --cov-report term-missing tests/
@@ -58,7 +58,7 @@ jobs:
         python -m pip install mypy
     - name: Install dependencies
       run: |
-        python -m pip install .[tests]
+        python -m pip install .[tests] -f https://storage.googleapis.com/jax-releases/jax_releases.html
     - name: Type checking
       continue-on-error: true
       run: |