Support returning the metric tensor of ExpvalCost in tape mode

pennylane/qnodes/base.py

@@ -576,22 +576,32 @@ def _construct(self, args, kwargs):
         self.queue = []  #: list[Operation]: applied operations
         self.obs_queue = []  #: list[Observable]: applied observables
 
+        tape_mode = qml.tape_mode_active()
+        if tape_mode:
+            qml.disable_tape()
+
+        try:
         # set up the context for Operator entry
-        with self:
-            try:
-                # generate the program queue by executing the quantum circuit function
-                if self.mutable:
-                    # it's ok to directly pass auxiliary arguments since the circuit is re-constructed each time
-                    # (positional args must be replaced because parameter-shift differentiation requires Variables)
-                    res = self.func(*self.arg_vars, **kwargs)
-                else:
-                    # TODO: Maybe we should only convert the kwarg_vars that were actually given
-                    res = self.func(*self.arg_vars, **self.kwarg_vars)
-            except:
-                # The qfunc call may have failed because the user supplied bad parameters, which is why we must wipe the created Variables.
-                self.arg_vars = None
-                self.kwarg_vars = None
-                raise
+            with self:
+                try:
+                    # generate the program queue by executing the quantum circuit function
+                    if self.mutable:
+                        # it's ok to directly pass auxiliary arguments since the circuit is
+                        # re-constructed each time (positional args must be replaced because
+                        # parameter-shift differentiation requires Variables)
+                        res = self.func(*self.arg_vars, **kwargs)
+                    else:
+                        # TODO: Maybe we should only convert the kwarg_vars that were actually given
+                        res = self.func(*self.arg_vars, **self.kwarg_vars)
+                except:
+                    # The qfunc call may have failed because the user supplied bad parameters,
+                    # which is why we must wipe the created Variables.
+                    self.arg_vars = None
+                    self.kwarg_vars = None
+                    raise
+        finally:
+            if tape_mode:
+                qml.enable_tape()
 
         # check the validity of the circuit
         self._check_circuit(res)

pennylane/vqe/vqe.py

@@ -470,10 +470,29 @@ 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()
+
+        tape_mode = qml.tape_mode_active()
+        if tape_mode:
+            try:
+                qml.disable_tape()
+                @qml.qnode(device, interface=interface, diff_method=diff_method, **kwargs)
+                def qnode_for_metric_tensor_in_tape_mode(*qnode_args, _wires=wires, **qnode_kwargs):
+                    """The metric tensor cannot currently be calculated in tape-mode QNodes. As a
+                    short-term fix for VQECost, 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)
+                    return qml.expval(qml.PauliZ(0))
+                self._qnode_for_metric_tensor_in_tape_mode = qnode_for_metric_tensor_in_tape_mode
+            finally:
+                qml.enable_tape()
+
         self._optimize = optimize
 
         if self._optimize:
-            if not qml.tape_mode_active():
+            if not tape_mode:
                 raise ValueError(
                     "Observable optimization is only supported in tape mode. Tape "
                     "mode can be enabled with the command:\n"
@@ -482,8 +501,6 @@ def __init__(
 
             obs_groupings, coeffs_groupings = qml.grouping.group_observables(observables, coeffs)
 
-            wires = device.wires.tolist()
-
             @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"""
@@ -523,12 +540,11 @@ def metric_tensor(self, args, kwargs=None, diag_approx=False, only_construct=Fal
         Returns:
             array[float]: metric tensor
         """
-        if self._optimize:
-            raise ValueError(
-                "Evaluation of the metric tensor is not supported when using "
-                "optimized observables. Set the argument optimize=False to obtain "
-                "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
             )
+
         # all the qnodes share the same ansatz so we select the first
         return self.qnodes.qnodes[0].metric_tensor(
             args=args, kwargs=kwargs, diag_approx=diag_approx, only_construct=only_construct
@@ -546,6 +562,6 @@ class VQECost(ExpvalCost):
     def __init__(self, *args, **kwargs):
         warnings.warn(
             "Use of VQECost is deprecated and should be replaced with ExpvalCost",
-            DeprecationWarning,
+            DeprecationWarning, 2
         )
         super().__init__(*args, **kwargs)

tests/qnodes/test_qnode_base.py

@@ -1500,3 +1500,29 @@ def circuit(a, b, c, d, *args):
         # the user will evaluate the QNode with.
         node.set_trainable_args({0, 1, 6})
         assert node.get_trainable_args() == {0, 1, 6}
+
+
+def test_old_qnode_in_tape_mode():
+    """Test that the old QNode can still be evaluated when running in tape mode"""
+
+    # tape mode should not be active so that we can use the old QNode
+    assert not qml.tape_mode_active()
+
+    try:
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev)
+        def f(x):
+            qml.RX(x, wires=0)
+            qml.RY(0.4, wires=1)
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0))
+
+        qml.enable_tape()
+        f(0.4)
+
+        # check that tape mode is turned on again after evaluating the old QNode
+        assert qml.tape_mode_active()
+
+    finally:  # always make sure we turn off tape mode to prevent disrupting the other tests
+        qml.disable_tape()

tests/test_vqe.py

@@ -849,19 +849,58 @@ def test_optimize_grad_tf(self, tf_support):
 
         assert np.allclose(dc, big_hamiltonian_grad)
 
-    def test_metric_tensor(self):
+    def test_metric_tensor_tape_mode(self):
         """Test that an error is raised if the metric tensor is requested in optimize=True mode."""
-        if not qml.tape_mode_active():
+        # if not not qml.tape_mode_active():
+        #     pytest.skip("This test is only intended for tape mode")
+
+        # dev = qml.device("default.qubit", wires=4)
+        # hamiltonian = big_hamiltonian
+        #
+        # def ansatz(a, b, c, d):
+        #     qml.RX(a, wires=0)
+        #     qml.RX(b, wires=0)
+        #     qml.RX(c, wires=0)
+        #     qml.RX(d, wires=0)
+        #
+        # cost = qml.ExpvalCost(ansatz, hamiltonian, dev)
+        #
+        # args = qml.init.strong_ent_layers_uniform(2, 4, seed=1967)
+        # cost.metric_tensor([1, 2, 3, 4], only_construct=True)
+
+        if not not qml.tape_mode_active():
             pytest.skip("This test is only intended for tape mode")
 
-        dev = qml.device("default.qubit", wires=4)
-        hamiltonian = big_hamiltonian
-
-        cost = qml.ExpvalCost(qml.templates.StronglyEntanglingLayers, hamiltonian, dev, optimize=True)
+        def my_template(params, wires, **kwargs):
+            qml.RX(params[0], wires=wires[0])
+            qml.RX(params[1], wires=wires[1])
+            qml.CNOT(wires=wires)
 
-        with pytest.raises(ValueError, match="Evaluation of the metric tensor is not supported"):
-            cost.metric_tensor(None)
+        obs_list = [qml.PauliX(0) @ qml.PauliZ(1), qml.PauliZ(0) @ qml.PauliX(1)]
 
+        dev = qml.device("default.qubit", wires=2)
+        qnodes = qml.map(my_template, obs_list, dev, measure="expval")
+
+        params = [0.54, 0.12]
+        print(qnodes(params))
+        print(qnodes[0](params))
+        print(qnodes[0].metric_tensor(params))
+
+
+
+        # dev = qml.device("default.qubit", wires=2)
+        # assert not qml.tape_mode_active()
+        #
+        # def ansatz(a, b, wires):
+        #     qml.RX(a, wires=0)
+        #     qml.RX(b, wires=1)
+        #     qml.CNOT(wires=[0, 1])
+        #
+        # qnode = qml.map(ansatz, [qml.PauliZ(0) @ qml.PauliZ(1)], dev)
+        #
+        # # mt = qnode[0].metric_tensor([1, 2])
+        # # print(mt)
+        # print(qnode(0.1, 0.2))
 
 @pytest.mark.usefixtures("tape_mode")
 class TestAutogradInterface: