Support multiple devices with ExpvalCost (PennyLaneAI#927)

* Fix multiple devices in ExpvalCost in tape mode

* Extend test to non-tape mode

* Fix typo

* Apply suggestions

* Applied suggestions

* Fix CI

* Update pennylane/vqe/vqe.py

Co-authored-by: Josh Izaac <josh146@gmail.com>

Co-authored-by: Josh Izaac <josh146@gmail.com>

pennylane/vqe/vqe.py

@@ -16,6 +16,7 @@
 computations using PennyLane.
 """
 # pylint: disable=too-many-arguments, too-few-public-methods
+from collections import Sequence
 import itertools
 import warnings
 
@@ -470,21 +471,26 @@ def __init__(
         """QNodeCollection: The QNodes to be evaluated. Each QNode corresponds to the expectation
         value of each observable term after applying the circuit ansatz."""
 
-        wires = device.wires.tolist()
+        self._multiple_devices = isinstance(device, Sequence)
+        """Bool: Records if multiple devices are input"""
 
         tape_mode = qml.tape_mode_active()
         if tape_mode:
+
+            d = device[0] if self._multiple_devices else device
+            w = d.wires.tolist()
+
             try:
                 qml.disable_tape()
 
-                @qml.qnode(device, interface=interface, diff_method=diff_method, **kwargs)
+                @qml.qnode(d, interface=interface, diff_method=diff_method, **kwargs)
                 def qnode_for_metric_tensor_in_tape_mode(*qnode_args, **qnode_kwargs):
                     """The metric tensor cannot currently be calculated in tape-mode QNodes. As a
                     short-term fix for ExpvalCost, we create a non-tape mode QNode just for
                     calculation of the metric tensor. In doing so, we reintroduce the same
                     restrictions of the old QNode but allow users to access new functionality
                     such as measurement grouping and batch execution of the gradient."""
-                    ansatz(*qnode_args, wires=wires, **qnode_kwargs)
+                    ansatz(*qnode_args, wires=w, **qnode_kwargs)
                     return qml.expval(qml.PauliZ(0))
 
                 self._qnode_for_metric_tensor_in_tape_mode = qnode_for_metric_tensor_in_tape_mode
@@ -501,12 +507,15 @@ def qnode_for_metric_tensor_in_tape_mode(*qnode_args, **qnode_kwargs):
                     "qml.enable_tape()"
                 )
 
+            if self._multiple_devices:
+                raise ValueError("Using multiple devices is not supported when optimize=True")
+
             obs_groupings, coeffs_groupings = qml.grouping.group_observables(observables, coeffs)
 
             @qml.qnode(device, interface=interface, diff_method=diff_method, **kwargs)
             def circuit(*qnode_args, obs, **qnode_kwargs):
                 """Converting ansatz into a full circuit including measurements"""
-                ansatz(*qnode_args, wires=wires, **qnode_kwargs)
+                ansatz(*qnode_args, wires=w, **qnode_kwargs)
                 return [qml.expval(o) for o in obs]
 
             def cost_fn(*qnode_args, **qnode_kwargs):
@@ -542,6 +551,12 @@ def metric_tensor(self, args, kwargs=None, diag_approx=False, only_construct=Fal
         Returns:
             array[float]: metric tensor
         """
+        if self._multiple_devices:
+            warnings.warn(
+                "ExpvalCost was instantiated with multiple devices. Only the first device "
+                "will be used to evaluate the metric tensor."
+            )
+
         if qml.tape_mode_active():
             return self._qnode_for_metric_tensor_in_tape_mode.metric_tensor(
                 args=args, kwargs=kwargs, diag_approx=diag_approx, only_construct=only_construct

tests/test_vqe.py

@@ -19,6 +19,7 @@
 
 import pennylane as qml
 from pennylane.wires import Wires
+from pennylane.devices import DefaultQubit, DefaultMixed
 
 try:
     import torch
@@ -886,6 +887,44 @@ def circuit(params):
         mt2 = circuit.metric_tensor([p])
         assert np.allclose(mt, mt2)
 
+    def test_multiple_devices(self, mocker):
+        """Test that passing multiple devices to ExpvalCost works correctly"""
+
+        dev = [qml.device("default.qubit", wires=2), qml.device("default.mixed", wires=2)]
+        spy = mocker.spy(DefaultQubit, "apply")
+        spy2 = mocker.spy(DefaultMixed, "apply")
+
+        obs = [qml.PauliZ(0), qml.PauliZ(1)]
+        h = qml.Hamiltonian([1, 1], obs)
+
+        qnodes = qml.ExpvalCost(qml.templates.BasicEntanglerLayers, h, dev)
+        w = qml.init.basic_entangler_layers_uniform(3, 2, seed=1967)
+
+        res = qnodes(w)
+
+        spy.assert_called_once()
+        spy2.assert_called_once()
+
+        mapped = qml.map(qml.templates.BasicEntanglerLayers, obs, dev)
+        exp = sum(mapped(w))
+
+        assert np.allclose(res, exp)
+
+        with pytest.warns(UserWarning, match="ExpvalCost was instantiated with multiple devices."):
+            qnodes.metric_tensor([w])
+
+    def test_multiple_devices_opt_true(self):
+        """Test if a ValueError is raised when multiple devices are passed when optimize=True."""
+        if not qml.tape_mode_active():
+            pytest.skip("This test is only intended for tape mode")
+
+        dev = [qml.device("default.qubit", wires=2), qml.device("default.qubit", wires=2)]
+
+        h = qml.Hamiltonian([1, 1], [qml.PauliZ(0), qml.PauliZ(1)])
+
+        with pytest.raises(ValueError, match="Using multiple devices is not supported when"):
+            qml.ExpvalCost(qml.templates.StronglyEntanglingLayers, h, dev, optimize=True)
+
 
 @pytest.mark.usefixtures("tape_mode")
 class TestAutogradInterface: