BugFix: Controlled.generator correctly projects with control values (P…

…ennyLaneAI#5068)

**Context:**

`Controlled.generator` always projects onto the $\ket{1}$ space of the
control wire without considering control values.

**Description of the Change:**

Now `Controlled.generator` correctly projects the operator based on the
control value specified

**Benefits:**

Fixes a bug

---------

Co-authored-by: Christina Lee <[email protected]>

doc/releases/changelog-dev.md

@@ -116,6 +116,8 @@
 * `StatePrep` operations expanded onto more wires are now compatible with backprop.
   [(#5028)](https://github.com/PennyLaneAI/pennylane/pull/5028)
 
+* The return value of `Controlled.generator` now contains a projector that projects onto the correct subspace based on the control value specified.
+  [(#5068)](https://github.com/PennyLaneAI/pennylane/pull/5068)
 
 <h3>Contributors ✍️</h3>
 

pennylane/ops/functions/generator.py

@@ -21,7 +21,7 @@
 import numpy as np
 
 import pennylane as qml
-from pennylane.ops import Hamiltonian, SProd, Sum
+from pennylane.ops import Hamiltonian, SProd, Prod, Sum
 
 
 def _generator_hamiltonian(gen, op):
@@ -46,6 +46,7 @@ def _generator_hamiltonian(gen, op):
     return H
 
 
+# pylint: disable=no-member
 def _generator_prefactor(gen):
     r"""Return the generator as ```(obs, prefactor)`` representing
     :math:`G=p \hat{O}`, where
@@ -58,6 +59,9 @@ def _generator_prefactor(gen):
 
     prefactor = 1.0
 
+    if isinstance(gen, Prod):
+        gen = qml.simplify(gen)
+
     if isinstance(gen, Hamiltonian):
         gen = qml.dot(gen.coeffs, gen.ops)  # convert to Sum
 
@@ -73,6 +77,7 @@ def _generator_prefactor(gen):
             prefactor = abs_coeffs[0]
             coeffs = [c / prefactor for c in coeffs]
             return qml.dot(coeffs, ops), prefactor
+
     elif isinstance(gen, SProd):
         return gen.base, gen.scalar
 

pennylane/ops/op_math/controlled.py

@@ -556,8 +556,14 @@ def has_generator(self):
 
     def generator(self):
         sub_gen = self.base.generator()
-        proj_tensor = operation.Tensor(*(qml.Projector([1], wires=w) for w in self.control_wires))
-        return 1.0 * proj_tensor @ sub_gen
+        projectors = (
+            qml.Projector([val], wires=w) for val, w in zip(self.control_values, self.control_wires)
+        )
+
+        if qml.operation.active_new_opmath():
+            return qml.prod(*projectors, sub_gen)
+
+        return 1.0 * operation.Tensor(*projectors) @ sub_gen
 
     @property
     def has_adjoint(self):

tests/ops/op_math/test_controlled.py

@@ -451,24 +451,38 @@ def test_has_generator_false(self):
 
         assert op.has_generator is False
 
-    def test_generator(self):
+    @pytest.mark.parametrize("use_new_op_math", [True, False])
+    def test_generator(self, use_new_op_math):
         """Test that the generator is a tensor product of projectors and the base's generator."""
 
+        if use_new_op_math:
+            qml.operation.enable_new_opmath()
+
         base = qml.RZ(-0.123, wires="a")
-        op = Controlled(base, ("b", "c"))
+        control_values = [0, 1]
+        op = Controlled(base, ("b", "c"), control_values=control_values)
 
         base_gen, base_gen_coeff = qml.generator(base, format="prefactor")
         gen_tensor, gen_coeff = qml.generator(op, format="prefactor")
 
         assert base_gen_coeff == gen_coeff
 
-        for wire, ob in zip(op.control_wires, gen_tensor.operands):
-            assert isinstance(ob, qml.Projector)
-            assert ob.data == ([1],)
-            assert ob.wires == qml.wires.Wires(wire)
+        for wire, val in zip(op.control_wires, control_values):
+            ob = list(op for op in gen_tensor.operands if op.wires == qml.wires.Wires(wire))
+            assert len(ob) == 1
+            assert ob[0].data == ([val],)
+
+        ob = list(op for op in gen_tensor.operands if op.wires == base.wires)
+        assert len(ob) == 1
+        assert ob[0].__class__ is base_gen.__class__
+
+        expected = qml.exp(op.generator(), 1j * op.data[0])
+        assert qml.math.allclose(
+            expected.matrix(wire_order=["a", "b", "c"]), op.matrix(wire_order=["a", "b", "c"])
+        )
 
-        assert gen_tensor.operands[-1].__class__ is base_gen.__class__
-        assert gen_tensor.operands[-1].wires == base_gen.wires
+        if use_new_op_math:
+            qml.operation.disable_new_opmath()
 
     def test_diagonalizing_gates(self):
         """Test that the Controlled diagonalizing gates is the same as the base diagonalizing gates."""