Fix Plugin Test Matrix stable/stable tests (#1019)

**Context:** The plugin test matrix fails because of a recent pin to jaxlib version. **Description of the Change:** We'll follow the pennylane and Catalyst pins. [stable/stable](https://github.com/PennyLaneAI/pennylane-lightning/actions/runs/12201421220) **Benefits:** Tests will work. **Possible Drawbacks:** **Related GitHub Issues:** [sc-79942] --------- Co-authored-by: Ali Asadi <10773383+maliasadi@users.noreply.github.com>
PennyLaneAI · Dec 6, 2024 · 9eb6123 · 9eb6123
1 parent 5448c8a
commit 9eb6123
Show file tree

Hide file tree

Showing 6 changed files with 7 additions and 6 deletions.
diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
@@ -83,6 +83,9 @@
 
 ### Bug fixes
 
+* Pin `jax[cpu]==0.4.28` for compatibility with PennyLane and Catalyst.
+  [(#1019)](https://github.com/PennyLaneAI/pennylane-lightning/pull/1019)
+
 * Fix Lightning Kokkos editable mode path.
   [(#1010)](https://github.com/PennyLaneAI/pennylane-lightning/pull/1010)
 

diff --git a/.github/workflows/tests_lkcpu_python.yml b/.github/workflows/tests_lkcpu_python.yml
@@ -235,7 +235,6 @@ jobs:
       - name: Install ML libraries for interfaces
         run: |
           python -m pip install --upgrade torch==$TORCH_VERSION -f https://download.pytorch.org/whl/cpu/torch_stable.html
-          python -m pip install --upgrade "jax[cpu]"  # This also installs jaxlib
           python -m pip install --upgrade tensorflow~=$TF_VERSION keras~=$TF_VERSION
 
       - name: Run PennyLane-Lightning unit tests

diff --git a/.github/workflows/tests_lkcuda_python.yml b/.github/workflows/tests_lkcuda_python.yml
@@ -225,7 +225,7 @@ jobs:
           cd pennylane
           git checkout master
           python -m pip uninstall -y pennylane && python -m pip install . -vv --no-deps
-  
+
       - name: Switch to release build of PennyLane
         if: inputs.pennylane-version == 'release'
         run: |
@@ -243,7 +243,6 @@ jobs:
       - name: Install ML libraries for interfaces
         run: |
           python -m pip install --upgrade torch==$TORCH_VERSION -f https://download.pytorch.org/whl/cpu/torch_stable.html
-          python -m pip install --upgrade "jax[cpu]"  # This also installs jaxlib
           python -m pip install --upgrade tensorflow~=$TF_VERSION keras~=$TF_VERSION
 
       - name: Install backend device
@@ -258,7 +257,7 @@ jobs:
           PL_BACKEND=${{ matrix.pl_backend }} python scripts/configure_pyproject_toml.py || true
           PL_BACKEND=${{ matrix.pl_backend }} CMAKE_ARGS="-DCMAKE_PREFIX_PATH=${{ github.workspace }}/Kokkos" \
           python -m pip install . -vv
-  
+
       - name: Run PennyLane-Lightning unit tests
         if: ${{ matrix.pl_backend != 'all'}}
         env:

diff --git a/.github/workflows/tests_lqcpu_python.yml b/.github/workflows/tests_lqcpu_python.yml
@@ -183,7 +183,6 @@ jobs:
       - name: Install ML libraries for interfaces
         run: |
           python -m pip install --upgrade torch==$TORCH_VERSION -f https://download.pytorch.org/whl/cpu/torch_stable.html
-          python -m pip install --upgrade "jax[cpu]"  # This also installs jaxlib
           python -m pip install --upgrade tensorflow~=$TF_VERSION keras~=$TF_VERSION
 
       - name: Switch to stable tag of Lightning

diff --git a/pennylane_lightning/core/_version.py b/pennylane_lightning/core/_version.py
@@ -16,4 +16,4 @@
    Version number (major.minor.patch[-label])
 """
 
-__version__ = "0.40.0-dev31"
+__version__ = "0.40.0-dev32"
diff --git a/requirements-dev.txt b/requirements-dev.txt
@@ -15,6 +15,7 @@ clang-format~=16.0
 isort==5.13.2
 click==8.0.4
 cmake
+jax[cpu]==0.4.28
 custatevec-cu12
 cutensornet-cu12
 pylint==2.7.4