Repository: PennyLaneAI/pennylane
Pull request #5068 (2024-01-15 20:34:48)

BugFix: Controlled.generator correctly projects with control values
**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


codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5068?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Comparison is base [(`c01cb4b`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/c01cb4b96f758a9f1a0c4e422115829b80e113cb?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.67% compared to head [(`bfdd1c1`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5068?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.66%.
> Report is 14 commits behind head on master.


<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #5068      +/-   ##
==========================================
- Coverage   99.67%   99.66%   -0.02%     
==========================================
  Files         394      394              
  Lines       35670    35426     -244     
==========================================
- Hits        35554    35307     -247     
- Misses        116      119       +3     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5068?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).

astralcai: > We should also add some tests to ensure this definition of the generator agrees with the `qml.exp(op.generator(), 1j*op.data[0])` convention for a generator. `op` and that exponential operator should have the same form.
> 
> We could potentially even add this check to `assert_valid`:
> 
> https://github.com/PennyLaneAI/pennylane/blob/master/pennylane/ops/functions/assert_valid.py
> 
> to test generators more generally, but for now we should just check those two forms have the same matrix.

I added a quick check for this in the test case, is checking matrix equivalence the proper way of doing this? 

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 05c35fd0b..bcae207fa 100644
--- a/doc/releases/changelog-dev.md
+++ b/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>
 
diff --git a/pennylane/ops/functions/generator.py b/pennylane/ops/functions/generator.py
index 61cce2c67..a00cdeecc 100644
--- a/pennylane/ops/functions/generator.py
+++ b/pennylane/ops/functions/generator.py
@@ -21,7 +21,7 @@ import warnings
 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
 
diff --git a/pennylane/ops/op_math/controlled.py b/pennylane/ops/op_math/controlled.py
index 422d12489..766a683a7 100644
--- a/pennylane/ops/op_math/controlled.py
+++ b/pennylane/ops/op_math/controlled.py
@@ -556,8 +556,14 @@ class Controlled(SymbolicOp):
 
     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):


Test patch:
diff --git a/tests/ops/op_math/test_controlled.py b/tests/ops/op_math/test_controlled.py
index c5fa4d579..590a87e96 100644
--- a/tests/ops/op_math/test_controlled.py
+++ b/tests/ops/op_math/test_controlled.py
@@ -451,24 +451,38 @@ class TestMiscMethods:
 
         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."""