Merge branch 'fermi-commute' of github.com:willjmax/pennylane into fe…

…rmi-commute

.github/workflows/install_deps/action.yml

@@ -15,7 +15,7 @@ inputs:
   jax_version:
     description: The version of JAX to install for any job that requires JAX
     required: false
-    default: '0.4.23'
+    default: '0.4.28'
   install_tensorflow:
     description: Indicate if TensorFlow should be installed or not
     required: false
@@ -86,7 +86,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

doc/development/deprecations.rst

@@ -39,6 +39,12 @@ Pending deprecations
   - Deprecated in v0.37
   - Will be removed in v0.39
 
+* The ``QubitStateVector`` template is deprecated.
+  Instead, use ``StatePrep``.
+
+  - Deprecated in v0.39
+  - Will be removed in v0.40
+
 New operator arithmetic deprecations
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 

doc/releases/changelog-dev.md

@@ -10,6 +10,9 @@
   [(#6061)](https://github.com/PennyLaneAI/pennylane/pull/6061)
   [(#6258)](https://github.com/PennyLaneAI/pennylane/pull/6258)
 
+* PennyLane is now compatible with Jax 0.4.28.
+  [(#6255)](https://github.com/PennyLaneAI/pennylane/pull/6255)
+
 * `qml.qchem.excitations` now optionally returns fermionic operators.
   [(#6171)](https://github.com/PennyLaneAI/pennylane/pull/6171)
 
@@ -54,6 +57,9 @@
   unique representation of the object.
   [(#6167)](https://github.com/PennyLaneAI/pennylane/pull/6167)
 
+* A `ReferenceQubit` is introduced for testing purposes and as a reference for future plugin development.
+  [(#6181)](https://github.com/PennyLaneAI/pennylane/pull/6181)
+
 * The `to_mat` methods for `FermiWord` and `FermiSentence` now optionally return
   a sparse matrix.
   [(#6173)](https://github.com/PennyLaneAI/pennylane/pull/6173)
@@ -90,6 +96,10 @@
 
 <h3>Deprecations 👋</h3>
 
+* The ``QubitStateVector`` template is deprecated.
+   Instead, use ``StatePrep``.
+   [(#6172)](https://github.com/PennyLaneAI/pennylane/pull/6172)
+
 * `Device`, `QubitDevice`, and `QutritDevice` will no longer be accessible via top-level import in v0.40.
   They will still be accessible as `qml.devices.LegacyDevice`, `qml.devices.QubitDevice`, and `qml.devices.QutritDevice`
   respectively.

pennylane/_version.py

@@ -16,4 +16,4 @@
 Version number (major.minor.patch[-label])
 """
 
-__version__ = "0.39.0-dev15"
+__version__ = "0.39.0-dev16"

pennylane/devices/__init__.py

@@ -37,6 +37,7 @@
     _qubit_device
     _qutrit_device
     null_qubit
+    reference_qubit
     tests
 
 Next generation devices
@@ -58,6 +59,7 @@
     DefaultQubit
     DefaultTensor
     NullQubit
+    ReferenceQubit
     DefaultQutritMixed
     LegacyDeviceFacade
 
@@ -160,6 +162,7 @@ def execute(self, circuits, execution_config = qml.devices.DefaultExecutionConfi
 from .default_clifford import DefaultClifford
 from .default_tensor import DefaultTensor
 from .null_qubit import NullQubit
+from .reference_qubit import ReferenceQubit
 from .default_qutrit import DefaultQutrit
 from .default_qutrit_mixed import DefaultQutritMixed
 from ._legacy_device import Device as LegacyDevice

pennylane/devices/reference_qubit.py

@@ -0,0 +1,154 @@
+# Copyright 2018-2024 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Contains the ReferenceQubit device, a minimal device that can be used for testing
+and plugin development purposes.
+"""
+
+import numpy as np
+
+import pennylane as qml
+
+from .device_api import Device
+from .execution_config import DefaultExecutionConfig
+from .modifiers import simulator_tracking, single_tape_support
+from .preprocess import decompose, validate_device_wires, validate_measurements
+
+
+def sample_state(state: np.ndarray, shots: int, seed=None):
+    """Generate samples from the provided state and number of shots."""
+
+    probs = np.imag(state) ** 2 + np.real(state) ** 2
+    basis_states = np.arange(len(probs))
+
+    num_wires = int(np.log2(len(probs)))
+
+    rng = np.random.default_rng(seed)
+    basis_samples = rng.choice(basis_states, shots, p=probs)
+
+    # convert basis state integers to array of booleans
+    bin_strings = (format(s, f"0{num_wires}b") for s in basis_samples)
+    return np.array([[int(val) for val in s] for s in bin_strings])
+
+
+def simulate(tape: qml.tape.QuantumTape, seed=None) -> qml.typing.Result:
+    """Simulate a tape and turn it into results.
+
+    Args:
+        tape (.QuantumTape): a representation of a circuit
+        seed (Any): A seed to use to control the generation of samples.
+
+    """
+    # 1) create the initial state
+    state = np.zeros(2 ** len(tape.wires))
+    state[0] = 1.0
+
+    # 2) apply all the operations
+    for op in tape.operations:
+        op_mat = op.matrix(wire_order=tape.wires)
+        state = qml.math.matmul(op_mat, state)
+
+    # 3) perform measurements
+    # note that shots are pulled from the tape, not from the device
+    if tape.shots:
+        samples = sample_state(state, shots=tape.shots.total_shots, seed=seed)
+        # Shot vector support
+        results = []
+        for lower, upper in tape.shots.bins():
+            sub_samples = samples[lower:upper]
+            results.append(
+                tuple(mp.process_samples(sub_samples, tape.wires) for mp in tape.measurements)
+            )
+        if len(tape.measurements) == 1:
+            results = [res[0] for res in results]
+        if not tape.shots.has_partitioned_shots:
+            results = results[0]
+        else:
+            results = tuple(results)
+    else:
+        results = tuple(mp.process_state(state, tape.wires) for mp in tape.measurements)
+        if len(tape.measurements) == 1:
+            results = results[0]
+
+    return results
+
+
+operations = frozenset({"PauliX", "PauliY", "PauliZ", "Hadamard", "CNOT", "CZ", "RX", "RY", "RZ"})
+
+
+def supports_operation(op: qml.operation.Operator) -> bool:
+    """This function used by preprocessing determines what operations
+    are natively supported by the device.
+
+    While in theory ``simulate`` can support any operation with a matrix, we limit the target
+    gate set for improved testing and reference purposes.
+
+    """
+    return getattr(op, "name", None) in operations
+
+
+@simulator_tracking  # update device.tracker with some relevant information
+@single_tape_support  # add support for device.execute(tape) in addition to device.execute((tape,))
+class ReferenceQubit(Device):
+    """A slimmed down numpy-based simulator for reference and testing purposes.
+
+    Args:
+        wires (int, Iterable[Number, str]): Number of wires present on the device, or iterable that
+            contains unique labels for the wires as numbers (i.e., ``[-1, 0, 2]``) or strings
+            (``['aux', 'q1', 'q2']``). Default ``None`` if not specified. While this device allows
+            for ``wires`` to be unspecified at construction time, other devices may make this argument
+            mandatory. Devices can also implement additional restrictions on the possible wires.
+        shots (int, Sequence[int], Sequence[Union[int, Sequence[int]]]): The default number of shots
+            to use in executions involving this device. Note that during execution, shots
+            are pulled from the circuit, not from the device.
+        seed (Union[str, None, int, array_like[int], SeedSequence, BitGenerator, Generator, jax.random.PRNGKey]): A
+            seed-like parameter matching that of ``seed`` for ``numpy.random.default_rng``. This is an optional
+            keyword argument added to follow recommend NumPy best practices. Other devices do not need
+            this parameter if it does not make sense for them.
+
+    """
+
+    name = "reference.qubit"
+
+    def __init__(self, wires=None, shots=None, seed=None):
+        super().__init__(wires=wires, shots=shots)
+
+        # seed and rng not necessary for a device, but part of recommended
+        # numpy practices to use a local random number generator
+        self._rng = np.random.default_rng(seed)
+
+    def preprocess(self, execution_config=DefaultExecutionConfig):
+
+        # Here we convert an arbitrary tape into one natively supported by the device
+        program = qml.transforms.core.TransformProgram()
+        program.add_transform(validate_device_wires, wires=self.wires, name="reference.qubit")
+        program.add_transform(qml.defer_measurements)
+        program.add_transform(qml.transforms.split_non_commuting)
+        program.add_transform(qml.transforms.diagonalize_measurements)
+        program.add_transform(
+            decompose,
+            stopping_condition=supports_operation,
+            skip_initial_state_prep=False,
+            name="reference.qubit",
+        )
+        program.add_transform(validate_measurements, name="reference.qubit")
+        program.add_transform(qml.transforms.broadcast_expand)
+
+        # no need to preprocess the config as the device does not support derivatives
+        return program, execution_config
+
+    def execute(self, circuits, execution_config=DefaultExecutionConfig):
+        for tape in circuits:
+            assert all(supports_operation(op) for op in tape.operations)
+        return tuple(simulate(tape, seed=self._rng) for tape in circuits)

pennylane/ops/qubit/state_preparation.py

@@ -15,6 +15,8 @@
 This submodule contains the discrete-variable quantum operations concerned
 with preparing a certain state on the device.
 """
+import warnings
+
 # pylint:disable=too-many-branches,abstract-method,arguments-differ,protected-access,no-member
 from typing import Optional
 
@@ -442,9 +444,19 @@ def _preprocess(state, wires, pad_with, normalize, validate_norm):
         return state
 
 
-# pylint: disable=missing-class-docstring
 class QubitStateVector(StatePrep):
-    pass  # QSV is still available
+    r"""
+    ``QubitStateVector`` is deprecated and will be removed in version 0.40. Instead, please use ``StatePrep``.
+    """
+
+    # pylint: disable=too-many-arguments
+    def __init__(self, state, wires, pad_with=None, normalize=False, validate_norm=True):
+        warnings.warn(
+            "QubitStateVector is deprecated and will be removed in version 0.40. "
+            "Instead, please use StatePrep.",
+            qml.PennyLaneDeprecationWarning,
+        )
+        super().__init__(state, wires, pad_with, normalize, validate_norm)
 
 
 class QubitDensityMatrix(Operation):

setup.py

@@ -52,6 +52,7 @@
             "default.qubit.legacy = pennylane.devices:DefaultQubitLegacy",
             "default.gaussian = pennylane.devices:DefaultGaussian",
             "default.mixed = pennylane.devices.default_mixed:DefaultMixed",
+            "reference.qubit = pennylane.devices.reference_qubit:ReferenceQubit",
             "null.qubit = pennylane.devices.null_qubit:NullQubit",
             "default.qutrit = pennylane.devices.default_qutrit:DefaultQutrit",
             "default.clifford = pennylane.devices.default_clifford:DefaultClifford",

tests/devices/default_qubit/test_default_qubit.py

@@ -1864,6 +1864,7 @@ def circ_expected():
         if use_jit:
             import jax
 
+            pytest.xfail(reason="'shots' cannot be a static_argname for 'jit' in JAX 0.4.28")
             circ_postselect = jax.jit(circ_postselect, static_argnames=["shots"])
 
         res = circ_postselect(param, shots=shots)
@@ -2051,6 +2052,7 @@ def circ():
         if use_jit:
             import jax
 
+            pytest.xfail(reason="'shots' cannot be a static_argname for 'jit' in JAX 0.4.28")
             circ = jax.jit(circ, static_argnames=["shots"])
 
         res = circ(shots=shots)

tests/devices/qutrit_mixed/test_qutrit_mixed_measure.py

@@ -499,7 +499,7 @@ def test_jax_backprop(self, use_jit):
 
         x = jax.numpy.array(self.x, dtype=jax.numpy.float64)
         coeffs = (5.2, 6.7)
-        f = jax.jit(self.f, static_argnums=(1, 2, 3, 4)) if use_jit else self.f
+        f = jax.jit(self.f, static_argnums=(1, 2, 3)) if use_jit else self.f
 
         out = f(x, coeffs)
         expected_out = self.expected(x, coeffs)

tests/devices/test_default_qutrit_mixed.py

@@ -823,7 +823,7 @@ def test_jax_backprop(self, use_jit):
 
         x = jax.numpy.array(self.x, dtype=jax.numpy.float64)
         coeffs = (5.2, 6.7)
-        f = jax.jit(self.f, static_argnums=(1, 2, 3, 4)) if use_jit else self.f
+        f = jax.jit(self.f, static_argnums=(1, 2, 3)) if use_jit else self.f
 
         out = f(x, coeffs)
         expected_out = self.expected(x, coeffs)

tests/drawer/test_drawable_layers.py

@@ -196,7 +196,7 @@ def test_mid_measure_custom_wires(self):
         m1 = qml.measurements.MeasurementValue([mp1], lambda v: v)
 
         def teleport(state):
-            qml.QubitStateVector(state, wires=["A"])
+            qml.StatePrep(state, wires=["A"])
             qml.Hadamard(wires="a")
             qml.CNOT(wires=["a", "B"])
             qml.CNOT(wires=["A", "a"])