default.qubit returns the new DefaultQubit device

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

@@ -263,9 +263,9 @@ jobs:
     strategy:
       matrix:
         config:
-          - device: default.qubit
+          - device: default.qubit.legacy
             shots: None
-          - device: default.qubit
+          - device: default.qubit.legacy
             shots: 10000
           # - device: default.qubit.tf
           #   shots: None

doc/releases/changelog-dev.md

@@ -17,9 +17,10 @@
 * Quantum information transforms are updated to the new transform program system.
   [(#4569)](https://github.com/PennyLaneAI/pennylane/pull/4569)
 
-* `qml.devices.DefaultQubit` now implements the new device API. The old version of `default.qubit`
-  is still accessible via `qml.devices.DefaultQubitLegacy`, or via short name `default.qubit.legacy`.
+* `default.qubit` now implements the new device API. The old version of the device is still
+  accessible by the short name `default.qubit.legacy`, or directly via `qml.devices.DefaultQubitLegacy`.
   [(#4594)](https://github.com/PennyLaneAI/pennylane/pull/4594)
+  [(#4436)](https://github.com/PennyLaneAI/pennylane/pull/4436)
 
 <h3>Improvements 🛠</h3>
 
@@ -184,6 +185,12 @@
   which effectively just called `marginal_prob` with `np.abs(state) ** 2`.
   [(#4602)](https://github.com/PennyLaneAI/pennylane/pull/4602)
 
+* `default.qubit` now implements the new device API. If you initialize a device
+  with `qml.device("default.qubit")`, all functions and properties that were tied to the old
+  device API will no longer be on the device. The legacy version can still be accessed with
+  `qml.device("default.qubit.legacy", wires=n_wires)`.
+  [(#4436)](https://github.com/PennyLaneAI/pennylane/pull/4436)
+
 <h3>Deprecations 👋</h3>
 
 * The ``prep`` keyword argument in ``QuantumScript`` is deprecated and will be removed from `QuantumScript`.

pennylane/devices/default_qubit.py

@@ -146,7 +146,7 @@ def f(x):
     @property
     def name(self):
         """The name of the device."""
-        return "default.qubit.2"
+        return "default.qubit"
 
     # pylint:disable = too-many-arguments
     def __init__(
@@ -193,6 +193,7 @@ def supports_derivatives(
         if (
             execution_config.gradient_method == "backprop"
             and execution_config.device_options.get("max_workers", self._max_workers) is None
+            and execution_config.interface is not None
         ):
             return True
 

pennylane/devices/default_qubit_legacy.py

@@ -104,8 +104,8 @@ class DefaultQubitLegacy(QubitDevice):
             returns analytical results.
     """
 
-    name = "Default qubit PennyLane plugin"
-    short_name = "default.qubit"
+    name = "Default qubit PennyLane plugin (Legacy)"
+    short_name = "default.qubit.legacy"
     pennylane_requires = __version__
     version = __version__
     author = "Xanadu Inc."

pennylane/devices/device_api.py

@@ -41,9 +41,6 @@ class Device(abc.ABC):
     """A device driver that can control one or more backends. A backend can be either a physical
     Quantum Processing Unit or a virtual one such as a simulator.
 
-    Device drivers should be configured to run under :func:`~.enable_return`, the newer
-    return shape specification, as the old return shape specification is deprecated.
-
     Only the ``execute`` method must be defined to construct a device driver.
 
     .. details::

pennylane/devices/qubit/adjoint_jacobian.py

@@ -75,6 +75,8 @@ def adjoint_jacobian(tape: QuantumTape, state=None):
     param_number = len(tape.get_parameters(trainable_only=False, operations_only=True)) - 1
     trainable_param_number = len(tape.trainable_params) - 1
     for op in reversed(tape.operations[tape.num_preps :]):
+        if isinstance(op, qml.Snapshot):
+            continue
         adj_op = qml.adjoint(op)
         ket = apply_operation(adj_op, ket)
 

pennylane/math/single_dispatch.py

@@ -454,6 +454,7 @@ def _cond_tf(pred, true_fn, false_fn, args):
     "vander",
     lambda *args, **kwargs: _i("tf").experimental.numpy.vander(*args, **kwargs),
 )
+ar.register_function("tensorflow", "size", lambda x: _i("tf").size(x))
 
 
 # -------------------------------- Torch --------------------------------- #

pennylane/measurements/state.py

@@ -222,4 +222,4 @@ def process_state(self, state: Sequence[complex], wire_order: Wires):
         # pylint:disable=redefined-outer-name
         wire_map = dict(zip(wire_order, range(len(wire_order))))
         mapped_wires = [wire_map[w] for w in self.wires]
-        return qml.math.reduce_statevector(state, indices=mapped_wires, c_dtype=state.dtype)
+        return qml.math.reduce_statevector(state, indices=mapped_wires)

pennylane/optimize/spsa.py

@@ -262,11 +262,13 @@ def compute_grad(self, objective_fn, args, kwargs):
         yminus = objective_fn(*thetaminus, **kwargs)
         try:
             # pylint: disable=protected-access
-            shots = (
-                Shots(objective_fn.device._raw_shot_sequence)
-                if objective_fn.device.shot_vector is not None
-                else Shots(None)
-            )
+            dev_shots = objective_fn.device.shots
+            if isinstance(dev_shots, Shots):
+                shots = dev_shots if dev_shots.has_partitioned_shots else Shots(None)
+            elif objective_fn.device.shot_vector is not None:
+                shots = Shots(objective_fn.device._raw_shot_sequence)  # pragma: no cover
+            else:
+                shots = Shots(None)
             if np.prod(objective_fn.func(*args).shape(objective_fn.device, shots)) > 1:
                 raise ValueError(
                     "The objective function must be a scalar function for the gradient "

pennylane/qnode.py

@@ -705,7 +705,9 @@ def _validate_backprop_method(device, interface, shots=None):
             config = qml.devices.ExecutionConfig(gradient_method="backprop", interface=interface)
             if device.supports_derivatives(config):
                 return "backprop", {}, device
-            raise qml.QuantumFunctionError(f"Device {device.name} does not support backprop")
+            raise qml.QuantumFunctionError(
+                f"Device {device.name} does not support backprop with {config}"
+            )
 
         mapped_interface = INTERFACE_MAP.get(interface, interface)
 

setup.py

@@ -47,7 +47,7 @@
         # TODO: rename entry point 'pennylane.plugins' to 'pennylane.devices'.
         # This requires a rename in the setup file of all devices, and is best done during another refactor
         "pennylane.plugins": [
-            "default.qubit = pennylane.devices:DefaultQubitLegacy",
+            "default.qubit = pennylane.devices:DefaultQubit",
             "default.qubit.legacy = pennylane.devices:DefaultQubitLegacy",
             "default.gaussian = pennylane.devices:DefaultGaussian",
             "default.qubit.tf = pennylane.devices.default_qubit_tf:DefaultQubitTF",

tests/devices/experimental/test_default_qubit_2.py

@@ -27,7 +27,7 @@
 
 def test_name():
     """Tests the name of DefaultQubit."""
-    assert DefaultQubit().name == "default.qubit.2"
+    assert DefaultQubit().name == "default.qubit"
 
 
 def test_shots():
@@ -200,6 +200,56 @@ def test_tracking_resources(self):
         assert len(tracker.history["resources"]) == 1
         assert tracker.history["resources"][0] == expected_resources
 
+    def test_tracking_batched_execution(self):
+        """Test the number of times the device is executed over a QNode's
+        lifetime is tracked by the device's tracker."""
+
+        dev_1 = qml.device("default.qubit", wires=2)
+
+        def circuit_1(x, y):
+            qml.RX(x, wires=[0])
+            qml.RY(y, wires=[1])
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1))
+
+        node_1 = qml.QNode(circuit_1, dev_1)
+        num_evals_1 = 10
+
+        with qml.Tracker(dev_1, persistent=True) as tracker1:
+            for _ in range(num_evals_1):
+                node_1(0.432, np.array([0.12, 0.5, 3.2]))
+        assert tracker1.totals["executions"] == num_evals_1
+
+        # test a second instance of a default qubit device
+        dev_2 = qml.device("default.qubit", wires=2)
+
+        def circuit_2(x):
+            qml.RX(x, wires=[0])
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1))
+
+        node_2 = qml.QNode(circuit_2, dev_2)
+        num_evals_2 = 5
+
+        with qml.Tracker(dev_2) as tracker2:
+            for _ in range(num_evals_2):
+                node_2(np.array([0.432, 0.61, 8.2]))
+        assert tracker2.totals["executions"] == num_evals_2
+
+        # test a new circuit on an existing instance of a qubit device
+        def circuit_3(y):
+            qml.RY(y, wires=[1])
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1))
+
+        node_3 = qml.QNode(circuit_3, dev_1)
+        num_evals_3 = 7
+
+        with tracker1:
+            for _ in range(num_evals_3):
+                node_3(np.array([0.12, 1.214]))
+        assert tracker1.totals["executions"] == num_evals_1 + num_evals_3
+
 
 # pylint: disable=too-few-public-methods
 class TestPreprocessing:
@@ -302,7 +352,7 @@ def test_supports_backprop(self):
         assert dev.supports_jvp() is True
         assert dev.supports_vjp() is True
 
-        config = ExecutionConfig(gradient_method="backprop")
+        config = ExecutionConfig(gradient_method="backprop", interface="auto")
         assert dev.supports_derivatives(config) is True
         assert dev.supports_jvp(config) is True
         assert dev.supports_vjp(config) is True
@@ -317,6 +367,11 @@ def test_supports_backprop(self):
         assert dev.supports_jvp(config) is False
         assert dev.supports_vjp(config) is False
 
+        config = ExecutionConfig(gradient_method="backprop", interface=None)
+        assert dev.supports_derivatives(config) is False
+        assert dev.supports_jvp(config) is False
+        assert dev.supports_vjp(config) is False
+
     def test_supports_adjoint(self):
         """Test that DefaultQubit says that it supports adjoint differentiation."""
         dev = DefaultQubit()

tests/devices/test_default_qubit_autograd.py

@@ -487,7 +487,7 @@ def circuit(a, b):
 
         def cost(a, b):
             prob_wire_1 = circuit(a, b)
-            return prob_wire_1[1] - prob_wire_1[0]
+            return prob_wire_1[1] - prob_wire_1[0]  # pylint:disable=unsubscriptable-object
 
         res = cost(a, b)
         expected = -np.cos(a) * np.cos(b)
@@ -513,7 +513,7 @@ def circuit(a, b):
 
         def cost(a, b):
             prob_wire_1 = circuit(a, b)
-            return prob_wire_1[:, 1] - prob_wire_1[:, 0]
+            return prob_wire_1[:, 1] - prob_wire_1[:, 0]  # pylint:disable=unsubscriptable-object
 
         res = cost(a, b)
         expected = -np.cos(a) * np.cos(b)

tests/devices/test_default_qubit_legacy.py

@@ -643,7 +643,7 @@ def test_apply_errors_qubit_state_vector(self, qubit_device_2_wires):
         with pytest.raises(
             DeviceError,
             match="Operation StatePrep cannot be used after other Operations have already been applied "
-            "on a default.qubit device.",
+            "on a default.qubit.legacy device.",
         ):
             qubit_device_2_wires.reset()
             qubit_device_2_wires.apply(
@@ -664,7 +664,7 @@ def test_apply_errors_basis_state(self, qubit_device_2_wires):
         with pytest.raises(
             DeviceError,
             match="Operation BasisState cannot be used after other Operations have already been applied "
-            "on a default.qubit device.",
+            "on a default.qubit.legacy device.",
         ):
             qubit_device_2_wires.reset()
             qubit_device_2_wires.apply(
@@ -2091,7 +2091,7 @@ def test_apply_parametrized_evolution_raises_error(self):
         param_ev = qml.evolve(ParametrizedHamiltonian([1], [qml.PauliX(0)]))
         with pytest.raises(
             NotImplementedError,
-            match="The device default.qubit cannot execute a ParametrizedEvolution operation",
+            match="The device default.qubit.legacy cannot execute a ParametrizedEvolution operation",
         ):
             self.dev._apply_parametrized_evolution(state=self.state, operation=param_ev)
 

tests/devices/test_default_qubit_legacy_broadcasting.py

@@ -462,7 +462,7 @@ def test_apply_errors_qubit_state_vector_broadcasted(self, qubit_device_2_wires)
         with pytest.raises(
             DeviceError,
             match="Operation StatePrep cannot be used after other Operations have already been applied "
-            "on a default.qubit device.",
+            "on a default.qubit.legacy device.",
         ):
             qubit_device_2_wires.apply([qml.RZ(0.5, wires=[0]), vec])
 
@@ -491,7 +491,7 @@ def test_apply_errors_basis_state_broadcasted(self, qubit_device_2_wires):
         with pytest.raises(
             DeviceError,
             match="Operation BasisState cannot be used after other Operations have already been applied "
-            "on a default.qubit device.",
+            "on a default.qubit.legacy device.",
         ):
             qubit_device_2_wires.apply([vec])
 

tests/devices/test_default_qubit_tf.py

@@ -1873,7 +1873,7 @@ def circuit(a, b):
             # get the probability of wire 1
             prob_wire_1 = circuit(a, b)
             # compute Prob(|1>_1) - Prob(|0>_1)
-            res = prob_wire_1[1] - prob_wire_1[0]
+            res = prob_wire_1[1] - prob_wire_1[0]  # pylint:disable=unsubscriptable-object
 
         expected = -tf.cos(a) * tf.cos(b)
         assert np.allclose(res, expected, atol=tol, rtol=0)
@@ -1900,7 +1900,7 @@ def circuit(a, b):
             # get the probability of wire 1
             prob_wire_1 = circuit(a, b)
             # compute Prob(|1>_1) - Prob(|0>_1)
-            res = prob_wire_1[:, 1] - prob_wire_1[:, 0]
+            res = prob_wire_1[:, 1] - prob_wire_1[:, 0]  # pylint:disable=unsubscriptable-object
 
         expected = -tf.cos(a) * tf.cos(b)
         assert np.allclose(res, expected, atol=tol, rtol=0)
@@ -1938,6 +1938,7 @@ def circuit(a, b):
             [-0.5 * np.sin(a) * (np.cos(b) + 1), 0.5 * np.sin(b) * (1 - np.cos(a))]
         )
 
+        # pylint:disable=no-member
         assert np.allclose(res.numpy(), expected_cost, atol=tol, rtol=0)
 
         res = tape.gradient(res, [a_tf, b_tf])
@@ -1971,6 +1972,7 @@ def circuit(a, b):
             [-0.5 * np.sin(a) * (np.cos(b) + 1), 0.5 * np.sin(b) * (1 - np.cos(a))]
         )
 
+        # pylint:disable=no-member
         assert np.allclose(res.numpy(), expected_cost, atol=tol, rtol=0)
 
         jac = tape.jacobian(res, [a_tf, b_tf])
@@ -2112,8 +2114,8 @@ def circuit(a):
         res = circuit(a)
 
         assert isinstance(res, tf.Tensor)
-        assert res.shape == (shots,)
-        assert set(res.numpy()) == {-1, 1}
+        assert res.shape == (shots,)  # pylint:disable=comparison-with-callable
+        assert set(res.numpy()) == {-1, 1}  # pylint:disable=no-member
 
     def test_estimating_marginal_probability(self, tol):
         """Test that the probability of a subset of wires is accurately estimated."""
@@ -2190,8 +2192,8 @@ def circuit(a):
         res = circuit(a)
 
         assert isinstance(res, tf.Tensor)
-        assert res.shape == (3, shots)
-        assert set(res.numpy().flat) == {-1, 1}
+        assert res.shape == (3, shots)  # pylint:disable=comparison-with-callable
+        assert set(res.numpy().flat) == {-1, 1}  # pylint:disable=no-member
 
     @pytest.mark.parametrize("batch_size", [2, 3])
     def test_estimating_marginal_probability_broadcasted(self, batch_size, tol):