Controlled operations rework Part 1

doc/releases/changelog-dev.md

@@ -165,9 +165,17 @@
   and the codecov check itself would never execute.
   [(#5101)](https://github.com/PennyLaneAI/pennylane/pull/5101)
 
+* `qml.ctrl` called on operators with custom controlled versions will return instances
+  of the custom class, and it will also flatten nested controlled operators to a single
+  multi-controlled operation. For `PauliX`, `CNOT`, `Toffoli`, and `MultiControlledX`,
+  calling `qml.ctrl` will always resolve to the best option in `CNOT`, `Toffoli`, or
+  `MultiControlledX` depending on the number of control wires and control values.
+  [(#5125)](https://github.com/PennyLaneAI/pennylane/pull/5125/)
+
 * `qml.Identity()` can be initialized without wires. Measuring it is currently not possible though.
   [(#5106)](https://github.com/PennyLaneAI/pennylane/pull/5106)
 
+
 <h4>Community contributions 🥳</h4>
 
 * The transform `split_non_commuting` now accepts measurements of type `probs`, `sample` and `counts` which accept both wires and observables.
@@ -203,7 +211,12 @@
   (with potentially negative eigenvalues) has been implemented.
   [(#5048)](https://github.com/PennyLaneAI/pennylane/pull/5048)
 
-* The decomposition of an operator created with calling `qml.ctrl` on a parametric operator (specifically `RX`, `RY`, `RZ`, `Rot`, `PhaseShift`) with a single control wire will now be the full decomposition instead of a single controlled gate. For example:
+* Controlled operators with a custom controlled version decomposes like how their
+  controlled counterpart decomposes, as opposed to decomposing into their controlled version.   
+  [(#5069)](https://github.com/PennyLaneAI/pennylane/pull/5069)
+  [(#5125)](https://github.com/PennyLaneAI/pennylane/pull/5125/)
+
+  For example:
   ```
   >>> qml.ctrl(qml.RX(0.123, wires=1), control=0).decomposition()
   [
@@ -215,7 +228,6 @@
     RZ(-1.5707963267948966, wires=[1])
   ]
   ```
-  [(#5069)](https://github.com/PennyLaneAI/pennylane/pull/5069)
 
 * `QuantumScript.is_sampled` and `QuantumScript.all_sampled` have been removed. Users should now
   validate these properties manually.

pennylane/devices/qubit/apply_operation.py

@@ -274,9 +274,7 @@ def apply_multicontrolledx(
         return _apply_operation_default(op, state, is_state_batched, debugger)
     ctrl_wires = [w + is_state_batched for w in op.control_wires]
     # apply x on all control wires with control value 0
-    roll_axes = [
-        w for val, w in zip(op.hyperparameters["control_values"], ctrl_wires) if val == "0"
-    ]
+    roll_axes = [w for val, w in zip(op.control_values, ctrl_wires) if val is False]
     for ax in roll_axes:
         state = math.roll(state, 1, ax)
 

pennylane/drawer/tape_mpl.py

@@ -107,8 +107,7 @@ def _(op: ops.Toffoli, drawer, layer, _):
 
 @_add_operation_to_drawer.register
 def _(op: ops.MultiControlledX, drawer, layer, _):
-    control_values = [(i == "1") for i in op.hyperparameters["control_values"]]
-    drawer.CNOT(layer, op.wires, control_values=control_values)
+    drawer.CNOT(layer, op.active_wires, control_values=op.control_values)
 
 
 @_add_operation_to_drawer.register

pennylane/ops/functions/bind_new_parameters.py

@@ -113,8 +113,9 @@ def bind_new_parameters_composite_op(op: CompositeOp, params: Sequence[TensorLik
 
 @bind_new_parameters.register(qml.CY)
 @bind_new_parameters.register(qml.CZ)
+@bind_new_parameters.register(qml.MultiControlledX)
 def bind_new_parameters_copy(
-    op: Union[qml.CY, qml.CZ], params: Sequence[TensorLike]
+    op: Union[qml.CY, qml.CZ, qml.MultiControlledX], params: Sequence[TensorLike]
 ):  # pylint:disable=unused-argument
     return copy.copy(op)
 

pennylane/ops/op_math/__init__.py

@@ -109,6 +109,7 @@
     CRZ,
     CY,
     CZ,
+    MultiControlledX,
 )
 from .decompositions import one_qubit_decomposition, two_qubit_decomposition, sk_decomposition
 from .evolution import Evolution
@@ -124,6 +125,7 @@
     "ControlledQubitUnitary",
     "CY",
     "CZ",
+    "MultiControlledX",
     "CRX",
     "CRY",
     "CRZ",

pennylane/ops/op_math/controlled.py

@@ -15,6 +15,7 @@
 This submodule defines the symbolic operation that indicates the control of an operator.
 """
 import warnings
+import functools
 from copy import copy
 from functools import wraps
 from inspect import signature
@@ -127,36 +128,41 @@ def cond_fn():
         ops_loader = available_eps[active_jit]["ops"].load()
         return ops_loader.ctrl(op, control, control_values=control_values, work_wires=work_wires)
 
-    custom_ops = {
-        (qml.PauliZ, 1): qml.CZ,
-        (qml.PauliY, 1): qml.CY,
-        (qml.PauliX, 1): qml.CNOT,
-        (qml.PauliX, 2): qml.Toffoli,
-        (qml.RX, 1): qml.CRX,
-        (qml.RY, 1): qml.CRY,
-        (qml.RZ, 1): qml.CRZ,
-        (qml.Rot, 1): qml.CRot,
-        (qml.PhaseShift, 1): qml.ControlledPhaseShift,
-    }
-    control_values = [control_values] if isinstance(control_values, (int, bool)) else control_values
     control = qml.wires.Wires(control)
-    custom_key = (type(op), len(control))
+    if isinstance(control_values, (int, bool)):
+        control_values = [control_values]
+    elif control_values is None:
+        control_values = [True] * len(control)
 
-    if custom_key in custom_ops and (control_values is None or all(control_values)):
-        qml.QueuingManager.remove(op)
-        return custom_ops[custom_key](*op.data, control + op.wires)
-    if isinstance(op, qml.PauliX):
+    ctrl_op = _try_wrap_in_custom_ctrl_op(
+        op, control, control_values=control_values, work_wires=work_wires
+    )
+    if ctrl_op is not None:
+        return ctrl_op
+
+    pauli_x_based_ctrl_ops = _get_pauli_x_based_ops()
+
+    # Special handling for PauliX-based controlled operations
+    if isinstance(op, pauli_x_based_ctrl_ops):
         qml.QueuingManager.remove(op)
-        control_string = (
-            None if control_values is None else "".join([str(int(v)) for v in control_values])
-        )
-        return qml.MultiControlledX(
-            wires=control + op.wires, control_values=control_string, work_wires=work_wires
+        return _handle_pauli_x_based_controlled_ops(op, control, control_values, work_wires)
+
+    # Flatten nested controlled operations to a multi-controlled operation for better
+    # decomposition algorithms. This includes special cases like CRX, CRot, etc.
+    if isinstance(op, Controlled):
+        work_wires = work_wires or []
+        return ctrl(
+            op.base,
+            control=control + op.control_wires,
+            control_values=control_values + op.control_values,
+            work_wires=work_wires + op.work_wires,
         )
+
     if isinstance(op, Operator):
         return Controlled(
             op, control_wires=control, control_values=control_values, work_wires=work_wires
         )
+
     if not callable(op):
         raise ValueError(
             f"The object {op} of type {type(op)} is not an Operator or callable. "
@@ -190,6 +196,115 @@ def wrapper(*args, **kwargs):
     return wrapper
 
 
+@functools.lru_cache()
+def _get_special_ops():
+    """Gets a list of special operations with custom controlled versions.
+
+    This is placed inside a function to avoid circular imports.
+
+    """
+
+    ops_with_custom_ctrl_ops = {
+        (qml.PauliZ, 1): qml.CZ,
+        (qml.PauliZ, 2): qml.CCZ,
+        (qml.PauliY, 1): qml.CY,
+        (qml.CZ, 1): qml.CCZ,
+        (qml.SWAP, 1): qml.CSWAP,
+        (qml.Hadamard, 1): qml.CH,
+        (qml.RX, 1): qml.CRX,
+        (qml.RY, 1): qml.CRY,
+        (qml.RZ, 1): qml.CRZ,
+        (qml.Rot, 1): qml.CRot,
+        (qml.PhaseShift, 1): qml.ControlledPhaseShift,
+    }
+    return ops_with_custom_ctrl_ops
+
+
+@functools.lru_cache()
+def _get_pauli_x_based_ops():
+    """Gets a list of pauli-x based operations
+
+    This is placed inside a function to avoid circular imports.
+
+    """
+    return qml.PauliX, qml.CNOT, qml.Toffoli, qml.MultiControlledX
+
+
+def _try_wrap_in_custom_ctrl_op(op, control, control_values=None, work_wires=None):
+    """Wraps a controlled operation in custom ControlledOp, returns None if not applicable."""
+
+    ops_with_custom_ctrl_ops = _get_special_ops()
+    custom_key = (type(op), len(control))
+
+    if custom_key in ops_with_custom_ctrl_ops and all(control_values):
+        qml.QueuingManager.remove(op)
+        return ops_with_custom_ctrl_ops[custom_key](*op.data, control + op.wires)
+
+    if isinstance(op, qml.QubitUnitary):
+        return qml.ControlledQubitUnitary(
+            op, control_wires=control, control_values=control_values, work_wires=work_wires
+        )
+
+    # A controlled ControlledPhaseShift should not be compressed to a multi controlled
+    # PhaseShift because the decomposition of PhaseShift contains a GlobalPhase that we
+    # do not have a controlled version of.
+    # TODO: remove this special case when we support ControlledGlobalPhase (sc-44933)
+    if isinstance(op, qml.ControlledPhaseShift):
+        return Controlled(
+            op, control_wires=control, control_values=control_values, work_wires=work_wires
+        )
+    # Similarly, compress the bottom levels of a multi-controlled PhaseShift to a
+    # ControlledPhaseShift if possible to avoid dealing with a controlled GlobalPhase
+    # during decomposition. This should also be removed in the future.
+    if isinstance(op, qml.PhaseShift) and control_values[-1]:
+        op = qml.ControlledPhaseShift(*op.data, wires=control[-1:] + op.wires)
+        return Controlled(
+            op,
+            control_wires=control[:-1],
+            control_values=control_values[:-1],
+            work_wires=work_wires,
+        )
+
+    return None
+
+
+def _handle_pauli_x_based_controlled_ops(op, control, control_values, work_wires):
+    """Handles PauliX-based controlled operations."""
+
+    op_map = {
+        (qml.PauliX, 1): qml.CNOT,
+        (qml.PauliX, 2): qml.Toffoli,
+        (qml.CNOT, 1): qml.Toffoli,
+    }
+
+    custom_key = (type(op), len(control))
+    if custom_key in op_map and all(control_values):
+        qml.QueuingManager.remove(op)
+        return op_map[custom_key](wires=control + op.wires)
+
+    if isinstance(op, qml.PauliX):
+        return qml.MultiControlledX(
+            wires=control + op.wires, control_values=control_values, work_wires=work_wires
+        )
+
+    # TODO: remove special handling of CNOT and Toffoli when they inherit from Controlled
+    if isinstance(op, qml.CNOT):
+        return qml.MultiControlledX(
+            wires=control + op.wires, control_values=control_values + [1], work_wires=work_wires
+        )
+    if isinstance(op, qml.Toffoli):
+        return qml.MultiControlledX(
+            wires=control + op.wires, control_values=control_values + [1, 1], work_wires=work_wires
+        )
+
+    work_wires = work_wires or []
+    return qml.MultiControlledX(
+        wires=control + op.wires,
+        control_values=control_values + op.control_values,
+        work_wires=work_wires + op.work_wires,
+    )
+
+
 # pylint: disable=too-many-arguments, too-many-public-methods
 class Controlled(SymbolicOp):
     """Symbolic operator denoting a controlled operator.
@@ -531,7 +646,7 @@ def has_decomposition(self):
             return True
         if len(self.control_wires) == 1 and hasattr(self.base, "_controlled"):
             return True
-        if isinstance(self.base, qml.PauliX):
+        if isinstance(self.base, _get_pauli_x_based_ops()):
             return True
         if _is_single_qubit_special_unitary(self.base):
             return True
@@ -631,21 +746,57 @@ def _is_single_qubit_special_unitary(op):
     return qmlmath.allclose(det, 1)
 
 
-# pylint: disable=protected-access
-def _decompose_no_control_values(op: "operation.Operator") -> List["operation.Operator"]:
-    """Provides a decomposition without considering control values. Returns None if
-    no decomposition.
-    """
+def _decompose_pauli_x_based_no_control_values(op: Controlled):
+    """Decomposes a PauliX-based operation"""
+
+    if isinstance(op.base, qml.PauliX) and len(op.control_wires) == 1:
+        return [qml.CNOT(wires=op.active_wires)]
+
+    if isinstance(op.base, qml.PauliX) and len(op.control_wires) == 2:
+        return qml.Toffoli.compute_decomposition(wires=op.active_wires)
+
+    if isinstance(op.base, qml.CNOT) and len(op.control_wires) == 1:
+        return qml.Toffoli.compute_decomposition(wires=op.active_wires)
+
+    return qml.MultiControlledX.compute_decomposition(
+        wires=op.active_wires,
+        work_wires=op.work_wires,
+    )
+
+
+def _decompose_custom_ops(op: Controlled) -> List["operation.Operator"]:
+    """Custom handling for decomposing a controlled operation"""
+
+    pauli_x_based_ctrl_ops = _get_pauli_x_based_ops()
+    ops_with_custom_ctrl_ops = _get_special_ops()
+
+    custom_key = (type(op.base), len(op.control_wires))
+    if custom_key in ops_with_custom_ctrl_ops:
+        custom_op_cls = ops_with_custom_ctrl_ops[custom_key]
+        return custom_op_cls.compute_decomposition(*op.data, op.active_wires)
+    if type(op.base) in pauli_x_based_ctrl_ops:
+        # has some special case handling of its own for further decomposition
+        return _decompose_pauli_x_based_no_control_values(op)
+
+    # TODO: will be removed in the second part of the controlled rework
     if len(op.control_wires) == 1 and hasattr(op.base, "_controlled"):
-        result = op.base._controlled(op.control_wires[0])
+        result = op.base._controlled(op.control_wires[0])  # pylint: disable=protected-access
         # disallow decomposing to itself
         # pylint: disable=unidiomatic-typecheck
         if type(result) != type(op):
             return [result]
         qml.QueuingManager.remove(result)
-    if isinstance(op.base, qml.PauliX):
-        # has some special case handling of its own for further decomposition
-        return [qml.MultiControlledX(wires=op.active_wires, work_wires=op.work_wires)]
+
+    return None
+
+
+def _decompose_no_control_values(op: Controlled) -> List["operation.Operator"]:
+    """Decompose without considering control values. Returns None if no decomposition."""
+
+    decomp = _decompose_custom_ops(op)
+    if decomp is not None:
+        return decomp
+
     if _is_single_qubit_special_unitary(op.base):
         if len(op.control_wires) >= 2 and qmlmath.get_interface(*op.data) == "numpy":
             return ctrl_decomp_bisect(op.base, op.control_wires)
@@ -663,7 +814,7 @@ def _decompose_no_control_values(op: "operation.Operator") -> List["operation.Op
             UserWarning,
         )
 
-    return [Controlled(newop, op.control_wires, work_wires=op.work_wires) for newop in base_decomp]
+    return [ctrl(newop, op.control_wires, work_wires=op.work_wires) for newop in base_decomp]
 
 
 class ControlledOp(Controlled, operation.Operation):