Repository: PennyLaneAI/pennylane
Pull request #1674 (2021-09-21 12:23:35)

Allow for operations to have gradient recipes that depend on the operation state
**Context:** Currently, operations cannot have gradient recipes (or gradient methods) that are dependent on the state of the operation. This is because the `grad_method`/`grad_recipe` validation occurs in `Operation.__init__` _prior_ to the operation attributes being set.

**Description of the Change:** The validation and construction logic is swapped. By placing the validation logic last, it is possible to write operations with state-dependent gradient methods/recipes.

For example:

```python
class RX(qml.RX):

    @property
    def grad_recipe(self):
        # The gradient is given by [f(2x) - f(0)] / (2 sin(x)), by subsituting
        # shift = x into the two term parameter-shift rule.
        x = self.data[0]
        c = 0.5 / np.sin(x)
        return ([[c, 0.0, 2 * x], [-c, 0.0, 0.0]],)

dev = qml.device("default.qubit", wires=2)

@qml.qnode(dev)
def circuit(x):
    RX(x, wires=0)
    return qml.expval(qml.PauliZ(0))

>>> circuit(0.5)
0.8775825618903726
>>> qml.grad(circuit)(0.5)
-0.479425538604203
```

We can also view the gradient tapes:

```python
>>> circuit.construct([0.5])
>>> tapes, fn = qml.gradients.param_shift(circuit.qtape)
>>> for t in tapes:
...     print(t.draw())
 0: ──RX(0)──┤ ⟨Z⟩
 0: ──RX(1)──┤ ⟨Z⟩
```

**Benefits:** More flexible gradient recipes.

**Possible Drawbacks:** There is still room to make the gradient recipe _more_ flexible. E.g.:

- The gradient recipe should be a function of the shift rule free parameters
- The gradient recipe should allow for changing the _structure_ of the circuit, and not simply be restricted to a linear combination of the same circuit with different parameters.

**Related GitHub Issues:** n/a



Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 2c676ad6b..683e7e513 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -263,6 +263,23 @@
 * ``qml.circuit_drawer.CircuitDrawer`` can accept a string for the ``charset`` keyword, instead of a ``CharSet`` object.
   [(#1640)](https://github.com/PennyLaneAI/pennylane/pull/1640)
 
+* Operations can now have gradient recipes that depend on the state of the operation.
+  [(#1674)](https://github.com/PennyLaneAI/pennylane/pull/1674)
+
+  For example, this allows for gradient recipes that are parameter dependent:
+
+  ```python
+  class RX(qml.RX):
+
+      @property
+      def grad_recipe(self):
+          # The gradient is given by [f(2x) - f(0)] / (2 sin(x)), by subsituting
+          # shift = x into the two term parameter-shift rule.
+          x = self.data[0]
+          c = 0.5 / np.sin(x)
+          return ([[c, 0.0, 2 * x], [-c, 0.0, 0.0]],)
+  ```
+
 <h3>Breaking changes</h3>
 
 - The `QNode.metric_tensor` method has been deprecated, and will be removed in an upcoming release.
diff --git a/pennylane/operation.py b/pennylane/operation.py
index afaa349ac..0c2a41dff 100644
--- a/pennylane/operation.py
+++ b/pennylane/operation.py
@@ -780,6 +780,7 @@ class Operation(Operator):
     def __init__(self, *params, wires=None, do_queue=True, id=None):
 
         self._inverse = False
+        super().__init__(*params, wires=wires, do_queue=do_queue, id=id)
 
         # check the grad_method validity
         if self.par_domain == "N":
@@ -804,8 +805,6 @@ class Operation(Operator):
         else:
             assert self.grad_recipe is None, "Gradient recipe is only used by the A method!"
 
-        super().__init__(*params, wires=wires, do_queue=do_queue, id=id)
-
 
 class DiagonalOperation(Operation):
     r"""Base class for diagonal quantum operations supported by a device.


Test patch:
diff --git a/tests/gradients/test_parameter_shift.py b/tests/gradients/test_parameter_shift.py
index d59d9e174..1b59d3e3a 100644
--- a/tests/gradients/test_parameter_shift.py
+++ b/tests/gradients/test_parameter_shift.py
@@ -263,6 +263,36 @@ class TestParamShift:
         assert np.allclose(j1, [exp, 0])
         assert np.allclose(j2, [0, exp])
 
+    def test_grad_recipe_parameter_dependent(self):
+        """Test that an operation with a gradient recipe that depends on
+        its instantiated parameter values works correctly within the parameter
+        shift rule"""
+
+        class RX(qml.RX):
+            @property
+            def grad_recipe(self):
+                # The gradient is given by [f(2x) - f(0)] / (2 sin(x)), by subsituting
+                # shift = x into the two term parameter-shift rule.
+                x = self.data[0]
+                c = 0.5 / np.sin(x)
+                return ([[c, 0.0, 2 * x], [-c, 0.0, 0.0]],)
+
+        x = np.array(0.654, requires_grad=True)
+        dev = qml.device("default.qubit", wires=2)
+
+        with qml.tape.JacobianTape() as tape:
+            RX(x, wires=0)
+            qml.expval(qml.PauliZ(0))
+
+        tapes, fn = qml.gradients.param_shift(tape)
+
+        assert len(tapes) == 2
+        assert tapes[0].operations[0].data[0] == 0
+        assert tapes[1].operations[0].data[0] == 2 * x
+
+        grad = fn(dev.batch_execute(tapes))
+        assert np.allclose(grad, -np.sin(x))
+
 
 class TestParameterShiftRule:
     """Tests for the parameter shift implementation"""
diff --git a/tests/test_operation.py b/tests/test_operation.py
index 0dae786fe..42a93e2e7 100644
--- a/tests/test_operation.py
+++ b/tests/test_operation.py
@@ -301,7 +301,7 @@ class TestOperationConstruction:
 
         class DummyOp(qml.operation.CVOperation):
             r"""Dummy custom operation"""
-            num_wires = 1
+            num_wires = 2
             num_params = 1
             par_domain = "R"
             grad_method = "A"
@@ -317,7 +317,7 @@ class TestOperationConstruction:
 
         class DummyOp(qml.operation.Operation):
             r"""Dummy custom operation"""
-            num_wires = 1
+            num_wires = 2
             num_params = 1
             par_domain = "N"
             grad_method = "A"
@@ -333,7 +333,7 @@ class TestOperationConstruction:
 
         class DummyOp(qml.operation.Operation):
             r"""Dummy custom operation"""
-            num_wires = 1
+            num_wires = 2
             num_params = 1
             par_domain = "A"
             grad_method = "A"
@@ -349,7 +349,7 @@ class TestOperationConstruction:
 
         class DummyOp(qml.operation.Operation):
             r"""Dummy custom operation"""
-            num_wires = 1
+            num_wires = 2
             num_params = 1
             par_domain = "R"
             grad_method = "F"
@@ -358,6 +358,26 @@ class TestOperationConstruction:
         with pytest.raises(AssertionError, match="Gradient recipe is only used by the A method"):
             DummyOp(0.5, wires=[0, 1])
 
+    def test_grad_recipe_parameter_dependent(self):
+        """Test that an operation with a gradient recipe that depends on
+        its instantiated parameter values works correctly"""
+
+        class DummyOp(qml.operation.Operation):
+            r"""Dummy custom operation"""
+            num_wires = 1
+            num_params = 1
+            par_domain = "R"
+            grad_method = "A"
+
+            @property
+            def grad_recipe(self):
+                x = self.data[0]
+                return ([[1.0, 1.0, x], [1.0, 0.0, -x]],)
+
+        x = 0.654
+        op = DummyOp(x, wires=0)
+        assert op.grad_recipe == ([[1.0, 1.0, x], [1.0, 0.0, -x]],)
+
     def test_no_wires_passed(self):
         """Test exception raised if no wires are passed"""