PennyLane is compatible with JAX 0.4.28

.github/workflows/install_deps/action.yml

@@ -8,30 +8,40 @@ inputs:
     description: The version of Python to use in order to run unit tests
     required: false
     default: '3.10'
+  install_numpy_2:
+    description: Indicate if numpy 2.0+ should be installed or not
+    required: false
+    type: boolean
+    default: false
+  numpy_2_version:
+    description: The version of numpy 2.0 to use in order to run unit tests
+    required: false
+    type: string
+    default: '2.0.0'
   install_jax:
     description: Indicate if JAX should be installed or not
     required: false
     default: 'true'
   jax_version:
     description: The version of JAX to install for any job that requires JAX
     required: false
-    default: 0.4.23
+    default: '0.4.26'
   install_tensorflow:
     description: Indicate if TensorFlow should be installed or not
     required: false
     default: 'true'
   tensorflow_version:
     description: The version of TensorFlow to install for any job that requires TensorFlow
     required: false
-    default: 2.16.0
+    default: '2.16.0'
   install_pytorch:
     description: Indicate if PyTorch should be installed or not
     required: false
     default: 'true'
   pytorch_version:
     description: The version of PyTorch to install for any job that requires PyTorch
     required: false
-    default: 2.3.0
+    default: '2.3.0'
   install_pennylane_lightning_master:
     description: Indicate if PennyLane-Lightning should be installed from the master branch
     required: false
@@ -67,6 +77,13 @@ runs:
         pip install -r requirements-ci.txt --upgrade
         pip install -r requirements-dev.txt --upgrade
 
+    - name: Install numpy 2.0 version
+      shell: bash
+      if: inputs.install_numpy_2 == 'true'
+      env:
+        NUMPY2_VERSION: ${{ inputs.numpy_2_version != '' && format('~={0}', inputs.numpy_2_version) || '' }}
+      run: pip install "numpy${{ env.NUMPY2_VERSION }}"
+
     - name: Install PyTorch
       shell: bash
       if: inputs.install_pytorch == 'true'
@@ -86,7 +103,7 @@ runs:
       if: inputs.install_jax == 'true'
       env:
         JAX_VERSION: ${{ inputs.jax_version != '' && format('=={0}', inputs.jax_version) || '' }}
-      run: pip install "jax${{ env.JAX_VERSION}}" "jaxlib${{ env.JAX_VERSION }}" scipy~=1.12.0
+      run: pip install "jax${{ env.JAX_VERSION}}" "jaxlib${{ env.JAX_VERSION }}"
 
     - name: Install additional PIP packages
       shell: bash

.github/workflows/interface-unit-tests.yml

@@ -152,14 +152,16 @@ jobs:
             fromJSON(needs.setup-ci-load.outputs.python-version).torch-tests
             || fromJSON(needs.setup-ci-load.outputs.python-version).default
            }}
+        numpy-2-install: [true, false]
     if: ${{ !contains(fromJSON(needs.setup-ci-load.outputs.jobs-to-skip), 'torch-tests') }}
     uses: ./.github/workflows/unit-test.yml
     with:
-      job_name: torch-tests (${{ matrix.python-version }})
+      job_name: torch-tests (${{ matrix.python-version }}, numpy 2.0 ${{ matrix.numpy-2-install }})
       branch: ${{ inputs.branch }}
-      coverage_artifact_name: core-interfaces-coverage-torch-${{ matrix.python-version }}
+      coverage_artifact_name: core-interfaces-coverage-torch-${{ matrix.python-version }}-${{ matrix.numpy-2-install }}
       python_version: ${{ matrix.python-version }}
       pipeline_mode: ${{ inputs.pipeline_mode }}
+      install_numpy_2: ${{ matrix.numpy-2-install }}
       install_jax: false
       install_tensorflow: false
       install_pytorch: true
@@ -185,14 +187,16 @@ jobs:
             fromJSON(needs.setup-ci-load.outputs.python-version).autograd-tests
             || fromJSON(needs.setup-ci-load.outputs.python-version).default
            }}
+        numpy-2-install: [true, false]    
     if: ${{ !contains(fromJSON(needs.setup-ci-load.outputs.jobs-to-skip), 'autograd-tests') }}
     uses: ./.github/workflows/unit-test.yml
     with:
-      job_name: autograd-tests (${{ matrix.python-version }})
+      job_name: autograd-tests (${{ matrix.python-version }}, numpy 2.0 ${{ matrix.numpy-2-install }})
       branch: ${{ inputs.branch }}
-      coverage_artifact_name: core-interfaces-coverage-autograd-${{ matrix.python-version }}
+      coverage_artifact_name: core-interfaces-coverage-autograd-${{ matrix.python-version }}-${{ matrix.numpy-2-install }}
       python_version: ${{ matrix.python-version }}
       pipeline_mode: ${{ inputs.pipeline_mode }}
+      install_numpy_2: ${{ matrix.numpy-2-install }}
       install_jax: false
       install_tensorflow: false
       install_pytorch: false
@@ -218,14 +222,16 @@ jobs:
             fromJSON(needs.setup-ci-load.outputs.python-version).tf-tests
             || fromJSON(needs.setup-ci-load.outputs.python-version).default
            }}
+        numpy-2-install: [true, false]    
     if: ${{ !contains(fromJSON(needs.setup-ci-load.outputs.jobs-to-skip), 'tf-tests') }}
     uses: ./.github/workflows/unit-test.yml
     with:
-      job_name: tf-tests (${{ matrix.group }}, ${{ matrix.python-version }})
+      job_name: tf-tests (${{ matrix.group }}, ${{ matrix.python-version }}, numpy 2.0 ${{ matrix.numpy-2-install }})
       branch: ${{ inputs.branch }}
-      coverage_artifact_name: core-interfaces-coverage-tf-${{ matrix.python-version }}-${{ matrix.group }}
+      coverage_artifact_name: core-interfaces-coverage-tf-${{ matrix.python-version }}-${{ matrix.group }}-${{ matrix.numpy-2-install }}
       python_version: ${{ matrix.python-version }}
       pipeline_mode: ${{ inputs.pipeline_mode }}
+      install_numpy_2: ${{ matrix.numpy-2-install }}
       install_jax: false
       install_tensorflow: true
       install_pytorch: false
@@ -256,14 +262,16 @@ jobs:
             fromJSON(needs.setup-ci-load.outputs.python-version).jax-tests
             || fromJSON(needs.setup-ci-load.outputs.python-version).default
            }}
+        numpy-2-install: [true, false]
     if: ${{ !contains(fromJSON(needs.setup-ci-load.outputs.jobs-to-skip), 'jax-tests') }}
     uses: ./.github/workflows/unit-test.yml
     with:
-      job_name: jax-tests (${{ matrix.group }}, ${{ matrix.python-version }})
+      job_name: jax-tests (${{ matrix.group }}, ${{ matrix.python-version }}, numpy 2.0 ${{ matrix.numpy-2-install }})
       branch: ${{ inputs.branch }}
-      coverage_artifact_name: core-interfaces-coverage-jax-${{ matrix.python-version }}-${{ matrix.group }}
+      coverage_artifact_name: core-interfaces-coverage-jax-${{ matrix.python-version }}-${{ matrix.group }}-${{ matrix.numpy-2-install }}
       python_version: ${{ matrix.python-version }}
       pipeline_mode: ${{ inputs.pipeline_mode }}
+      install_numpy_2: ${{ matrix.numpy-2-install }}
       install_jax: true
       install_tensorflow: false
       install_pytorch: false
@@ -294,14 +302,16 @@ jobs:
             fromJSON(needs.setup-ci-load.outputs.python-version).core-tests
             || fromJSON(needs.setup-ci-load.outputs.python-version).default
            }}
+        numpy-2-install: [true, false]    
     if: ${{ !contains(fromJSON(needs.setup-ci-load.outputs.jobs-to-skip), 'core-tests') }}
     uses: ./.github/workflows/unit-test.yml
     with:
-      job_name: core-tests (${{ matrix.group }}, ${{ matrix.python-version }})
+      job_name: core-tests (${{ matrix.group }}, ${{ matrix.python-version }}, numpy 2.0 ${{ matrix.numpy-2-install }})
       branch: ${{ inputs.branch }}
-      coverage_artifact_name: core-interfaces-coverage-core-${{ matrix.python-version }}-${{ matrix.group }}
+      coverage_artifact_name: core-interfaces-coverage-core-${{ matrix.python-version }}-${{ matrix.group }}-${{ matrix.numpy-2-install }}
       python_version: ${{ matrix.python-version }}
       pipeline_mode: ${{ inputs.pipeline_mode }}
+      install_numpy_2: ${{ matrix.numpy-2-install }}
       install_jax: false
       install_tensorflow: false
       install_pytorch: false
@@ -331,14 +341,16 @@ jobs:
             fromJSON(needs.setup-ci-load.outputs.python-version).all-interfaces-tests
             || fromJSON(needs.setup-ci-load.outputs.python-version).default
            }}
+        numpy-2-install: [true, false]
     if: ${{ !contains(fromJSON(needs.setup-ci-load.outputs.jobs-to-skip), 'all-interfaces-tests') }}
     uses: ./.github/workflows/unit-test.yml
     with:
-      job_name: all-interfaces-tests (${{ matrix.python-version }})
+      job_name: all-interfaces-tests (${{ matrix.python-version }}, numpy 2.0 ${{ matrix.numpy-2-install }})
       branch: ${{ inputs.branch }}
-      coverage_artifact_name: all-interfaces-coverage
+      coverage_artifact_name: all-interfaces-coverage-${{ matrix.numpy-2-install }}
       python_version: ${{ matrix.python-version }}
       pipeline_mode: ${{ inputs.pipeline_mode }}
+      install_numpy_2: ${{ matrix.numpy-2-install }}
       install_jax: true
       install_tensorflow: true
       install_pytorch: true
@@ -364,14 +376,16 @@ jobs:
             fromJSON(needs.setup-ci-load.outputs.python-version).external-libraries-tests
             || fromJSON(needs.setup-ci-load.outputs.python-version).default
            }}
+        numpy-2-install: [true, false]     
     if: ${{ !contains(fromJSON(needs.setup-ci-load.outputs.jobs-to-skip), 'external-libraries-tests') }}
     uses: ./.github/workflows/unit-test.yml
     with:
-      job_name: external-libraries-tests (${{ matrix.python-version }})
+      job_name: external-libraries-tests (${{ matrix.python-version }}, numpy 2.0 ${{ matrix.numpy-2-install }})
       branch: ${{ inputs.branch }}
-      coverage_artifact_name: external-libraries-tests-coverage
+      coverage_artifact_name: external-libraries-tests-coverage-${{ matrix.numpy-2-install }}
       python_version: ${{ matrix.python-version }}
       pipeline_mode: ${{ inputs.pipeline_mode }}
+      install_numpy_2: ${{ matrix.numpy-2-install }}
       install_jax: true
       install_tensorflow: true
       install_pytorch: false
@@ -403,14 +417,16 @@ jobs:
             fromJSON(needs.setup-ci-load.outputs.python-version).qcut-tests
             || fromJSON(needs.setup-ci-load.outputs.python-version).default
            }}
+        numpy-2-install: [true, false]    
     if: ${{ !contains(fromJSON(needs.setup-ci-load.outputs.jobs-to-skip), 'qcut-tests') }}
     uses: ./.github/workflows/unit-test.yml
     with:
-      job_name: qcut-tests (${{ matrix.python-version }})
+      job_name: qcut-tests (${{ matrix.python-version }}, numpy 2.0 ${{ matrix.numpy-2-install }})
       branch: ${{ inputs.branch }}
-      coverage_artifact_name: qcut-coverage
+      coverage_artifact_name: qcut-coverage-${{ matrix.numpy-2-install }}
       python_version: ${{ matrix.python-version }}
       pipeline_mode: ${{ inputs.pipeline_mode }}
+      install_numpy_2: ${{ matrix.numpy-2-install }}
       install_jax: true
       install_tensorflow: true
       install_pytorch: true
@@ -436,14 +452,16 @@ jobs:
             fromJSON(needs.setup-ci-load.outputs.python-version).qchem-tests
             || fromJSON(needs.setup-ci-load.outputs.python-version).default
            }}
+        numpy-2-install: [true, false]
     if: ${{ !contains(fromJSON(needs.setup-ci-load.outputs.jobs-to-skip), 'qchem-tests') }}
     uses: ./.github/workflows/unit-test.yml
     with:
-      job_name: qchem-tests (${{ matrix.python-version }})
+      job_name: qchem-tests (${{ matrix.python-version }}, numpy 2.0 ${{ matrix.numpy-2-install }})
       branch: ${{ inputs.branch }}
-      coverage_artifact_name: qchem-coverage
+      coverage_artifact_name: qchem-coverage-${{ matrix.numpy-2-install }}
       python_version: ${{ matrix.python-version }}
       pipeline_mode: ${{ inputs.pipeline_mode }}
+      install_numpy_2: ${{ matrix.numpy-2-install }}
       install_jax: false
       install_tensorflow: false
       install_pytorch: false
@@ -471,14 +489,16 @@ jobs:
             fromJSON(needs.setup-ci-load.outputs.python-version).gradients-tests
             || fromJSON(needs.setup-ci-load.outputs.python-version).default
            }}
+        numpy-2-install: [true, false]    
     if: ${{ !contains(fromJSON(needs.setup-ci-load.outputs.jobs-to-skip), 'gradients-tests') }}
     uses: ./.github/workflows/unit-test.yml
     with:
-      job_name: gradients-tests (${{ matrix.config.suite }}, ${{ matrix.python-version }})
+      job_name: gradients-tests (${{ matrix.config.suite }}, ${{ matrix.python-version }}, numpy 2.0 ${{ matrix.numpy-2-install }})
       branch: ${{ inputs.branch }}
-      coverage_artifact_name: gradients-${{ matrix.config.suite }}-coverage
+      coverage_artifact_name: gradients-${{ matrix.config.suite }}-coverage-${{ matrix.numpy-2-install }}
       python_version: ${{ matrix.python-version }}
       pipeline_mode: ${{ inputs.pipeline_mode }}
+      install_numpy_2: ${{ matrix.numpy-2-install }}
       install_jax: true
       install_tensorflow: true
       install_pytorch: true
@@ -503,14 +523,16 @@ jobs:
             fromJSON(needs.setup-ci-load.outputs.python-version).data-tests
             || fromJSON(needs.setup-ci-load.outputs.python-version).default
            }}
+        numpy-2-install: [true, false]
     if: ${{ !contains(fromJSON(needs.setup-ci-load.outputs.jobs-to-skip), 'data-tests') }}
     uses: ./.github/workflows/unit-test.yml
     with:
-      job_name: data-tests (${{ matrix.python-version }})
+      job_name: data-tests (${{ matrix.python-version }}, numpy 2.0 ${{ matrix.numpy-2-install }})
       branch: ${{ inputs.branch }}
-      coverage_artifact_name: data-coverage-${{ matrix.python-version }}
+      coverage_artifact_name: data-coverage-${{ matrix.python-version }}-${{ matrix.numpy-2-install }}
       python_version: ${{ matrix.python-version }}
       pipeline_mode: ${{ inputs.pipeline_mode }}
+      install_numpy_2: ${{ matrix.numpy-2-install }}
       install_jax: false
       install_tensorflow: false
       install_pytorch: false
@@ -543,14 +565,16 @@ jobs:
             fromJSON(needs.setup-ci-load.outputs.python-version).device-tests
             || fromJSON(needs.setup-ci-load.outputs.python-version).default
            }}
+        numpy-2-install: [true, false]    
     if: ${{ !contains(fromJSON(needs.setup-ci-load.outputs.jobs-to-skip), 'device-tests') }}
     uses: ./.github/workflows/unit-test.yml
     with:
-      job_name: device-tests (${{ matrix.config.device }}, ${{ matrix.config.shots }}, ${{ matrix.python-version }})
+      job_name: device-tests (${{ matrix.config.device }}, ${{ matrix.config.shots }}, ${{ matrix.python-version }}, numpy 2.0 ${{ matrix.numpy-2-install }})
       branch: ${{ inputs.branch }}
-      coverage_artifact_name: devices-coverage-${{ matrix.config.device }}-${{ matrix.config.shots }}
+      coverage_artifact_name: devices-coverage-${{ matrix.config.device }}-${{ matrix.config.shots }}-${{ matrix.numpy-2-install }}
       python_version: ${{ matrix.python-version }}
       pipeline_mode: ${{ inputs.pipeline_mode }}
+      install_numpy_2: ${{ matrix.numpy-2-install }}
       install_jax: ${{ !contains(matrix.config.skip_interface, 'jax') }}
       install_tensorflow: ${{ !contains(matrix.config.skip_interface, 'tf') }}
       install_pytorch: ${{ !contains(matrix.config.skip_interface, 'torch') }}

.github/workflows/unit-test.yml

@@ -24,6 +24,16 @@ on:
         required: false
         type: string
         default: '3.10'
+      install_numpy_2:
+        description: Indicate if numpy 2.0+ should be installed or not
+        required: false
+        type: boolean
+        default: false
+      numpy_2_version:
+        description: The version of numpy 2.0 to use in order to run unit tests
+        required: false
+        type: string
+        default: '2.0'
       pipeline_mode:
         description: The pipeline mode can be unit-tests, benchmarks, or reference-benchmark
         required: false
@@ -158,6 +168,8 @@ jobs:
         uses: ./.github/workflows/install_deps
         with:
           python_version: ${{ inputs.python_version }}
+          numpy_2_version: ${{ inputs.numpy_2_version }}
+          install_numpy_2: ${{ inputs.install_numpy_2 }}
           install_pytorch: ${{ inputs.install_pytorch }}
           install_tensorflow: ${{ inputs.install_tensorflow }}
           install_jax: ${{ inputs.install_jax }}

.gitignore

@@ -27,3 +27,5 @@ config.toml
 qml_debug.log
 datasets/*
 .benchmarks/*
+*.h5
+*.hdf5

pennylane/numpy/random.py

@@ -16,18 +16,19 @@
 it works with the PennyLane :class:`~.tensor` class.
 """
 
-from autograd.numpy import random as _random
+# isort: skip_file
 from numpy import __version__ as np_version
 from numpy.random import MT19937, PCG64, SFC64, Philox  # pylint: disable=unused-import
+from autograd.numpy import random as _random
 from packaging.specifiers import SpecifierSet
 from packaging.version import Version
 
 from .wrapper import tensor_wrapper, wrap_arrays
 
 wrap_arrays(_random.__dict__, globals())
 
-
 if Version(np_version) in SpecifierSet(">=0.17.0"):
+
     # pylint: disable=too-few-public-methods
     # pylint: disable=missing-class-docstring
     class Generator(_random.Generator):

requirements-ci.txt

@@ -1,5 +1,5 @@
-numpy
-scipy<1.13.0
+numpy~=1.26
+scipy<=1.13.0
 cvxpy
 cvxopt
 networkx

setup.py

@@ -21,7 +21,7 @@
     version = f.readlines()[-1].split()[-1].strip("\"'")
 
 requirements = [
-    "numpy<2.0",
+    "numpy<=2.0",
     "scipy",
     "networkx",
     "rustworkx>=0.14.0",