Repository: PennyLaneAI/pennylane (python)

Issue:
[BUG] `qml.ctrl` broken differentiability with parameter-shift/finite-diff
### Expected behavior

`qml.ctrl` supports differentiating over input parameters to the function to be controlled.

For example, for an operation being controlled, the following works:
```python
# Autograd
# ----
import pennylane as qml
from pennylane import numpy as np

dev = qml.device("default.qubit", wires=2)

@qml.qnode(dev, diff_method='parameter-shift')
def circuit(b):
    qml.QubitStateVector(np.array([1.0, -1.0], requires_grad=False) / np.sqrt(2), wires=0)
    qml.ctrl(qml.RY, control=0)(b, wires=[1])
    #qml.CRY(b, wires=[0, 1])
    return qml.expval(qml.PauliX(0))

b = np.array(0.123, requires_grad=True)

qml.grad(circuit)(b)
```

### Actual behavior

* With Autograd: error is raised
* With TF/Torch: the gradient is 0

### Additional information

*Swapping out to a dedicated controlled operation (e.g., `CRY`) solves the issue in specific cases.
* Backprop and adjoint work as expected.

### Source code

```shell
# TF
# ------
import pennylane as qml
from pennylane import numpy as np
import tensorflow as tf

dev = qml.device("default.qubit", wires=2)

@qml.qnode(dev, interface='tf', diff_method='parameter-shift')
def circuit(x):
    qml.QubitStateVector(np.array([1.0, -1.0], requires_grad=False) / np.sqrt(2), wires=0)
    qml.ctrl(qml.RY, control=0)(x, wires=[1])
    #qml.CRY(b, wires=[0, 1])
    return qml.expval(qml.PauliX(0))

b = tf.Variable(0.123, dtype=tf.float64)
with tf.GradientTape() as tape:
    res = circuit(b)

grad1 = tape.gradient(res, b)
grad1
```


### Tracebacks

```shell
~/xanadu/pennylane/pennylane/devices/default_qubit.py in apply(self, operations, rotations, **kwargs)
    215                 self._apply_basis_state(operation.parameters[0], operation.wires)
    216             else:
--> 217                 self._state = self._apply_operation(self._state, operation)
    218 
    219         # store the pre-rotated state

~/xanadu/pennylane/pennylane/devices/default_qubit.py in _apply_operation(self, state, operation)
    246         if len(wires) <= 2:
    247             # Einsum is faster for small gates
--> 248             return self._apply_unitary_einsum(state, matrix, wires)
    249 
    250         return self._apply_unitary(state, matrix, wires)

~/xanadu/pennylane/pennylane/devices/default_qubit.py in _apply_unitary_einsum(self, state, mat, wires)
    729         device_wires = self.map_wires(wires)
    730 
--> 731         mat = self._cast(self._reshape(mat, [2] * len(device_wires) * 2), dtype=self.C_DTYPE)
    732 
    733         # Tensor indices of the quantum state

TypeError: must be real number, not ArrayBox
```


### System information

```shell
Name: PennyLane
Version: 0.22.0.dev0
Summary: PennyLane is a Python quantum machine learning library by Xanadu Inc.
Home-page: https://github.com/XanaduAI/pennylane
Author: 
Author-email: 
License: Apache License 2.0
Location: /xanadu/pennylane
Requires: numpy, scipy, networkx, retworkx, autograd, toml, appdirs, semantic_version, autoray, cachetools, pennylane-lightning
Required-by: PennyLane-Lightning, amazon-braket-pennylane-plugin, PennyLane-Orquestra, pennylane-qulacs, PennyLane-Honeywell, PennyLane-AQT, PennyLane-PQ, PennyLane-Forest, PennyLane-qsharp, PennyLane-Qchem, PennyLane-IonQ, PennyLane-SF, PennyLane-Cirq, PennyLane-qiskit
Platform info:           Linux-5.13.0-30-generic-x86_64-with-glibc2.10
Python version:          3.8.5
Numpy version:           1.22.1
Scipy version:           1.7.3
Installed devices:
- lightning.qubit (PennyLane-Lightning-0.21.0)
- braket.aws.qubit (amazon-braket-pennylane-plugin-1.5.3)
- braket.local.qubit (amazon-braket-pennylane-plugin-1.5.3)
- orquestra.forest (PennyLane-Orquestra-0.15.0)
- orquestra.ibmq (PennyLane-Orquestra-0.15.0)
- orquestra.qiskit (PennyLane-Orquestra-0.15.0)
- orquestra.qulacs (PennyLane-Orquestra-0.15.0)
- qulacs.simulator (pennylane-qulacs-0.17.0.dev0)
- honeywell.hqs (PennyLane-Honeywell-0.16.0.dev0)
- aqt.noisy_sim (PennyLane-AQT-0.18.0)
- aqt.sim (PennyLane-AQT-0.18.0)
- projectq.classical (PennyLane-PQ-0.18.0.dev0)
- projectq.ibm (PennyLane-PQ-0.18.0.dev0)
- projectq.simulator (PennyLane-PQ-0.18.0.dev0)
- forest.numpy_wavefunction (PennyLane-Forest-0.18.0.dev0)
- forest.qvm (PennyLane-Forest-0.18.0.dev0)
- forest.wavefunction (PennyLane-Forest-0.18.0.dev0)
- microsoft.QuantumSimulator (PennyLane-qsharp-0.19.0)
- ionq.qpu (PennyLane-IonQ-0.17.0.dev0)
- ionq.simulator (PennyLane-IonQ-0.17.0.dev0)
- strawberryfields.fock (PennyLane-SF-0.20.0.dev0)
- strawberryfields.gaussian (PennyLane-SF-0.20.0.dev0)
- strawberryfields.gbs (PennyLane-SF-0.20.0.dev0)
- strawberryfields.remote (PennyLane-SF-0.20.0.dev0)
- strawberryfields.tf (PennyLane-SF-0.20.0.dev0)
- cirq.mixedsimulator (PennyLane-Cirq-0.20.0.dev0)
- cirq.pasqal (PennyLane-Cirq-0.20.0.dev0)
- cirq.qsim (PennyLane-Cirq-0.20.0.dev0)
- cirq.qsimh (PennyLane-Cirq-0.20.0.dev0)
- cirq.simulator (PennyLane-Cirq-0.20.0.dev0)
- qiskit.aer (PennyLane-qiskit-0.22.0.dev0)
- qiskit.basicaer (PennyLane-qiskit-0.22.0.dev0)
- qiskit.ibmq (PennyLane-qiskit-0.22.0.dev0)
- qiskit.ibmq.circuit_runner (PennyLane-qiskit-0.22.0.dev0)
- qiskit.ibmq.sampler (PennyLane-qiskit-0.22.0.dev0)
- default.gaussian (PennyLane-0.22.0.dev0)
- default.mixed (PennyLane-0.22.0.dev0)
- default.qubit (PennyLane-0.22.0.dev0)
- default.qubit.autograd (PennyLane-0.22.0.dev0)
- default.qubit.jax (PennyLane-0.22.0.dev0)
- default.qubit.tf (PennyLane-0.22.0.dev0)
- default.qubit.torch (PennyLane-0.22.0.dev0)
```


### Existing GitHub issues

- [X] I have searched existing GitHub issues to make sure the issue does not already exist.

Pull request description:
Fix bug to allow `qml.ctrl` to work with gradient transforms
**Context:** As per #2237, the `qml.ctrl` transform was not working with gradient transforms such as the parameter-shift rule. It does however work with backprop.

**Description of the Change:**

- Ensures that the tape parameters are up to date prior to expansion
- Sets `grad_method=None` to enforce decomposition if using the parameter-shift rule
- Adds tests

**Benefits:** Gradient transforms now work with `qml.ctrl`.

**Possible Drawbacks:** n/a

**Related GitHub Issues:** Fixes #2237


Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 60c36ef58..a0017ec47 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -156,6 +156,10 @@
 
 <h3>Bug fixes</h3>
 
+* The ``qml.ctrl`` transform now works correctly with gradient transforms
+  such as the parameter-shift rule.
+  [(#2238)](https://github.com/PennyLaneAI/pennylane/pull/2238)
+
 * Fixes a bug in which passing required arguments into operations as
   keyword arguments would throw an error because the documented call
   signature didn't match the function definition.
diff --git a/pennylane/transforms/control.py b/pennylane/transforms/control.py
index d09bda272..9d8398ca8 100644
--- a/pennylane/transforms/control.py
+++ b/pennylane/transforms/control.py
@@ -83,6 +83,7 @@ class ControlledOperation(Operation):
         control_wires: A wire or set of wires.
     """
 
+    grad_method = None
     num_wires = AnyWires
 
     def __init__(self, tape, control_wires, do_queue=True):
@@ -105,8 +106,11 @@ class ControlledOperation(Operation):
 
     def expand(self):
         tape = self.subtape
+        tape.set_parameters(self.data)
+
         for wire in self.control_wires:
             tape = expand_with_control(tape, wire)
+
         return tape
 
     def adjoint(self):


Test patch:
diff --git a/tests/transforms/test_control.py b/tests/transforms/test_control.py
index 882d6972b..3bc311cab 100644
--- a/tests/transforms/test_control.py
+++ b/tests/transforms/test_control.py
@@ -309,3 +309,97 @@ def test_controlled_template_and_operations():
     tape = expand_tape(tape)
     assert len(tape.operations) == 10
     assert all(o.name in {"CNOT", "CRX", "Toffoli"} for o in tape.operations)
+
+
+@pytest.mark.parametrize("diff_method", ["backprop", "parameter-shift", "finite-diff"])
+class TestDifferentiation:
+    """Tests for differentiation"""
+
+    def test_autograd(self, diff_method):
+        """Test differentiation using autograd"""
+        from pennylane import numpy as pnp
+
+        dev = qml.device("default.qubit", wires=2)
+        init_state = pnp.array([1.0, -1.0], requires_grad=False) / np.sqrt(2)
+
+        @qml.qnode(dev, diff_method=diff_method)
+        def circuit(b):
+            qml.QubitStateVector(init_state, wires=0)
+            qml.ctrl(qml.RY, control=0)(b, wires=[1])
+            return qml.expval(qml.PauliX(0))
+
+        b = pnp.array(0.123, requires_grad=True)
+        res = qml.grad(circuit)(b)
+        expected = np.sin(b / 2) / 2
+
+        assert np.allclose(res, expected)
+
+    def test_torch(self, diff_method):
+        """Test differentiation using torch"""
+        torch = pytest.importorskip("torch")
+
+        dev = qml.device("default.qubit", wires=2)
+        init_state = torch.tensor([1.0, -1.0], requires_grad=False) / np.sqrt(2)
+
+        @qml.qnode(dev, diff_method=diff_method, interface="torch")
+        def circuit(b):
+            qml.QubitStateVector(init_state, wires=0)
+            qml.ctrl(qml.RY, control=0)(b, wires=[1])
+            return qml.expval(qml.PauliX(0))
+
+        b = torch.tensor(0.123, requires_grad=True)
+        loss = circuit(b)
+        loss.backward()
+
+        res = b.grad.detach()
+        expected = np.sin(b.detach() / 2) / 2
+
+        assert np.allclose(res, expected)
+
+    @pytest.mark.parametrize("jax_interface", ["jax", "jax-python", "jax-jit"])
+    def test_jax(self, diff_method, jax_interface):
+        """Test differentiation using JAX"""
+
+        if diff_method == "backprop" and jax_interface != "jax":
+            pytest.skip("The backprop case only accepts interface='jax'")
+
+        jax = pytest.importorskip("jax")
+        jnp = jax.numpy
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev, diff_method=diff_method, interface=jax_interface)
+        def circuit(b):
+            init_state = np.array([1.0, -1.0]) / np.sqrt(2)
+            qml.QubitStateVector(init_state, wires=0)
+            qml.ctrl(qml.RY, control=0)(b, wires=[1])
+            return qml.expval(qml.PauliX(0))
+
+        b = jnp.array(0.123)
+        res = jax.grad(circuit)(b)
+        expected = np.sin(b / 2) / 2
+
+        assert np.allclose(res, expected)
+
+    def test_tf(self, diff_method):
+        """Test differentiation using TF"""
+        tf = pytest.importorskip("tensorflow")
+
+        dev = qml.device("default.qubit", wires=2)
+        init_state = tf.constant([1.0, -1.0], dtype=tf.complex128) / np.sqrt(2)
+
+        @qml.qnode(dev, diff_method=diff_method, interface="tf")
+        def circuit(b):
+            qml.QubitStateVector(init_state, wires=0)
+            qml.ctrl(qml.RY, control=0)(b, wires=[1])
+            return qml.expval(qml.PauliX(0))
+
+        b = tf.Variable(0.123, dtype=tf.float64)
+
+        with tf.GradientTape() as tape:
+            loss = circuit(b)
+
+        res = tape.gradient(loss, b)
+        expected = np.sin(b / 2) / 2
+
+        assert np.allclose(res, expected)