Add support for arbitrary linear combination gradient recipes

pennylane/operation.py

@@ -636,16 +636,30 @@ def get_parameter_shift(self, idx):
         """
         # get the gradient recipe for this parameter
         recipe = self.grad_recipe[idx]
-        multiplier, shift = (0.5, np.pi / 2) if recipe is None else recipe
+
+        # Default values
+        multiplier = 0.5
+        a = 1
+        shift = np.pi / 2
+
+        # We set the default recipe following:
+        # ∂f(x) = c1*f(a1*x+s1) + c2*f(a2*x+s2)
+        # where we express a positive and a negative shift by default
+        default_param_shift = [[multiplier, a, shift], [-multiplier, a, -shift]]
+        param_shift = default_param_shift if recipe is None else recipe
 
         # internal multiplier in the Variable
         var_mult = self.data[idx].mult
 
-        multiplier *= var_mult
-        if var_mult != 0:
-            # zero multiplier means the shift is unimportant
-            shift /= var_mult
-        return multiplier, shift
+        for elem in param_shift:
+
+            # Update the multiplier
+            elem[0] *= var_mult
+            if var_mult != 0:
+                # Update the shift
+                # zero multiplier means the shift is unimportant
+                elem[2] /= var_mult
+        return param_shift
 
     @property
     def generator(self):
@@ -1588,16 +1602,33 @@ def heisenberg_pd(self, idx):
         """
         # get the gradient recipe for this parameter
         recipe = self.grad_recipe[idx]
-        multiplier = 0.5 if recipe is None else recipe[0]
-        shift = np.pi / 2 if recipe is None else recipe[1]
+
+        # Default values
+        multiplier = 0.5
+        a = 1
+        shift = np.pi / 2
+
+        # We set the default recipe to as follows:
+        # ∂f(x) = 0.5*f(1*x+pi/2) - 0.5*f(1*x-pi/2)
+        default_param_shift = [[multiplier, a, shift], [-multiplier, a, -shift]]
+        param_shift = default_param_shift if recipe is None else recipe
+
+        pd = None # partial derivative of the transformation
 
         p = self.parameters
-        # evaluate the transform at the shifted parameter values
-        p[idx] += shift
-        U2 = self._heisenberg_rep(p)  # pylint: disable=assignment-from-none
-        p[idx] -= 2 * shift
-        U1 = self._heisenberg_rep(p)  # pylint: disable=assignment-from-none
-        return (U2 - U1) * multiplier  # partial derivative of the transformation
+
+        original_p_idx = p[idx]
+        for c, _a, s in param_shift:
+            # evaluate the transform at the shifted parameter values
+            p[idx] = _a * original_p_idx + s
+            U = self._heisenberg_rep(p)  # pylint: disable=assignment-from-none
+
+            if pd is None:
+                pd = c * U
+            else:
+                pd += c * U
+
+        return pd
 
     def heisenberg_tr(self, wires, inverse=False):
         r"""Heisenberg picture representation of the linear transformation carried

pennylane/ops/cv.py

@@ -138,7 +138,9 @@ class Squeezing(CVOperation):
     grad_method = "A"
 
     shift = 0.1
-    grad_recipe = [(0.5 / math.sinh(shift), shift), None]
+    multiplier = 0.5 / math.sinh(shift)
+    a = 1
+    grad_recipe = ([[multiplier, a, shift],[-multiplier, a, -shift]], None)
 
     @staticmethod
     def _heisenberg_rep(p):
@@ -180,7 +182,9 @@ class Displacement(CVOperation):
     grad_method = "A"
 
     shift = 0.1
-    grad_recipe = [(0.5 / shift, shift), None]
+    multiplier = 0.5 / shift
+    a = 1
+    grad_recipe = ([[multiplier, a, shift],[-multiplier, a, -shift]], None)
 
     @staticmethod
     def _heisenberg_rep(p):
@@ -278,8 +282,11 @@ class TwoModeSqueezing(CVOperation):
     par_domain = "R"
 
     grad_method = "A"
+
     shift = 0.1
-    grad_recipe = [(0.5 / math.sinh(shift), shift), None]
+    multiplier = 0.5 / math.sinh(shift)
+    a = 1
+    grad_recipe = ([[multiplier, a, shift],[-multiplier, a, -shift]], None)
 
     @staticmethod
     def _heisenberg_rep(p):
@@ -326,8 +333,11 @@ class QuadraticPhase(CVOperation):
     par_domain = "R"
 
     grad_method = "A"
+
     shift = 0.1
-    grad_recipe = [(0.5 / shift, shift)]
+    multiplier = 0.5 / shift
+    a = 1
+    grad_recipe = ([[multiplier, a, shift],[-multiplier, a, -shift]],)
 
     @staticmethod
     def _heisenberg_rep(p):
@@ -371,8 +381,11 @@ class ControlledAddition(CVOperation):
     par_domain = "R"
 
     grad_method = "A"
+
     shift = 0.1
-    grad_recipe = [(0.5 / shift, shift)]
+    multiplier = 0.5 / shift
+    a = 1
+    grad_recipe = ([[multiplier, a, shift],[-multiplier, a, -shift]],)
 
     @staticmethod
     def _heisenberg_rep(p):
@@ -417,8 +430,11 @@ class ControlledPhase(CVOperation):
     par_domain = "R"
 
     grad_method = "A"
+
     shift = 0.1
-    grad_recipe = [(0.5 / shift, shift)]
+    multiplier = 0.5 / shift
+    a = 1
+    grad_recipe = ([[multiplier, a, shift],[-multiplier, a, -shift]],)
 
     @staticmethod
     def _heisenberg_rep(p):

pennylane/qnodes/cv.py

@@ -181,20 +181,37 @@ def _pd_analytic(self, idx, args, kwargs, **options):
             temp_var.idx = n
             op.data[p_idx] = temp_var
 
-            multiplier, shift = op.get_parameter_shift(p_idx)
-
-            # shifted parameter values
-            shift_p1 = np.r_[args, args[idx] + shift]
-            shift_p2 = np.r_[args, args[idx] - shift]
+            param_shift = op.get_parameter_shift(p_idx)
 
             if not force_order2 and op.use_method != "B":
                 # basic parameter-shift method, for Gaussian CV gates
                 # succeeded by order-1 observables
-                # evaluate the circuit at two points with shifted parameter values
-                y2 = np.asarray(self.evaluate(shift_p1, kwargs))
-                y1 = np.asarray(self.evaluate(shift_p2, kwargs))
-                pd += (y2 - y1) * multiplier
+                # evaluate the circuit at multiple points with the linear
+                # combination of parameter values (in most cases at two points)
+                for multiplier, a, shift in param_shift:
+
+                    # shifted parameter values
+                    shift_p = np.r_[args, a * args[idx] + shift]
+
+                    term = multiplier * np.asarray(self.evaluate(shift_p, kwargs))
+                    pd += term
             else:
+                if len(param_shift) != 2:
+                    # TODO: check if more than two terms is supported
+                    raise NotImplementedError(
+                        f"Taking the analytic gradient for order-2 operators is\
+                        unsupported for {op} which contains a parameter with a\
+                        gradient recipe of more than two terms."
+                    )
+
+                # Get the shifts and the multipliers
+                pos_multiplier, a1, pos_shift = param_shift[0]
+                neg_multiplier, a2, neg_shift = param_shift[1]
+
+                # shifted parameter values
+                shift_p1 = np.r_[args, a1 * args[idx] + pos_shift]
+                shift_p2 = np.r_[args, a2 * args[idx] + neg_shift]
+
                 # order-2 parameter-shift method, for gaussian CV gates
                 # succeeded by order-2 observables
                 # evaluate transformed observables at the original parameter point
@@ -203,7 +220,7 @@ def _pd_analytic(self, idx, args, kwargs, **options):
                 Z2 = op.heisenberg_tr(self.device.wires)
                 self._set_variables(shift_p2, kwargs)
                 Z1 = op.heisenberg_tr(self.device.wires)
-                Z = (Z2 - Z1) * multiplier  # derivative of the operation
+                Z = pos_multiplier * Z2 + neg_multiplier * Z1  # derivative of the operation
 
                 unshifted_args = np.r_[args, args[idx]]
                 self._set_variables(unshifted_args, kwargs)

pennylane/qnodes/qubit.py

@@ -128,16 +128,18 @@ def _pd_analytic(self, idx, args, kwargs, **options):
             temp_var.idx = n
             op.data[p_idx] = temp_var
 
-            multiplier, shift = op.get_parameter_shift(p_idx)
+            param_shift = op.get_parameter_shift(p_idx)
 
-            # shifted parameter values
-            shift_p1 = np.r_[args, args[idx] + shift]
-            shift_p2 = np.r_[args, args[idx] - shift]
+            for multiplier, a, shift in param_shift:
 
-            # evaluate the circuit at two points with shifted parameter values
-            y2 = np.asarray(self.evaluate(shift_p1, kwargs))
-            y1 = np.asarray(self.evaluate(shift_p2, kwargs))
-            pd += (y2 - y1) * multiplier
+                # shifted parameter values
+                shift_p = np.r_[args, a * args[idx] + shift]
+
+                # evaluate the circuit at point with shifted parameter values
+                y = np.asarray(self.evaluate(shift_p, kwargs))
+
+                # add the contribution to the partial derivative
+                pd += multiplier * y
 
             # restore the original parameter
             op.data[p_idx] = orig

pennylane/tape/tapes/cv_param_shift.py

@@ -236,18 +236,30 @@ def parameter_shift_first_order(
         p_idx = self._par_info[t_idx]["p_idx"]
 
         recipe = op.grad_recipe[p_idx]
-        c, s = (0.5, np.pi / 2) if recipe is None else recipe
+
+        # Default values
+        multiplier = 0.5
+        a = 1
+        shift = np.pi / 2
+
+        # We set the default recipe to as follows:
+        # ∂f(x) = 0.5*f(1*x+pi/2) - 0.5*f(1*x-pi/2)
+        default_param_shift = [[multiplier, a, shift], [-multiplier, a, -shift]]
+        param_shift = default_param_shift if recipe is None else recipe
 
         shift = np.zeros_like(params)
-        shift[idx] = s
 
-        shifted_forward = self.copy(copy_operations=True, tape_cls=QuantumTape)
-        shifted_forward.set_parameters(params + shift)
+        coeffs = []
+        tapes = []
+        for c, _a, s in param_shift:
 
-        shifted_backward = self.copy(copy_operations=True, tape_cls=QuantumTape)
-        shifted_backward.set_parameters(params - shift)
+            shift[idx] = s
 
-        tapes = [shifted_forward, shifted_backward]
+            # shifted parameter values
+            shifted_tape = self.copy(copy_operations=True, tape_cls=QuantumTape)
+            shifted_tape.set_parameters(params + shift)
+            coeffs.append(c)
+            tapes.append(shifted_tape)
 
         def processing_fn(results):
             """Computes the gradient of the parameter at index idx via the
@@ -263,7 +275,14 @@ def processing_fn(results):
             shifted_forward = np.array(results[0])
             shifted_backward = np.array(results[1])
 
-            return c * (shifted_forward - shifted_backward)
+            stat = np.zeros_like(results[0])
+
+            for c, res in zip(coeffs, results):
+                shifted = np.array(res)
+                # stat += c * shifted
+                np.add(stat, c * shifted, out=stat, casting="unsafe")
+
+            return stat
 
         return tapes, processing_fn
 
@@ -291,21 +310,35 @@ def parameter_shift_second_order(self, idx, params, **options):
         dev_wires = options["dev_wires"]
 
         recipe = op.grad_recipe[p_idx]
-        c, s = (0.5, np.pi / 2) if recipe is None else recipe
+
+        # Default values
+        multiplier = 0.5
+        a = 1
+        shift = np.pi / 2
+
+        # We set the default recipe following:
+        # ∂f(x) = c1*f(a1*x+s1) + c2*f(a2*x+s2)
+        # where we express a positive and a negative shift by default
+        default_param_shift = [[multiplier, a, shift], [-multiplier, a, -shift]]
+        param_shift = default_param_shift if recipe is None else recipe
+
+        c1, a1, s1 = param_shift[0]
+        c2, a2, s2 = param_shift[1]
 
         shift = np.zeros_like(params)
-        shift[idx] = s
+        shift[idx] = s1
 
         # evaluate transformed observables at the original parameter point
         # first build the Heisenberg picture transformation matrix Z
-        self.set_parameters(params + shift)
+        self.set_parameters(a1 * params + shift)
         Z2 = op.heisenberg_tr(dev_wires)
 
-        self.set_parameters(params - shift)
+        shift[idx] = s2
+        self.set_parameters(a2 * params + shift)
         Z1 = op.heisenberg_tr(dev_wires)
 
         # derivative of the operation
-        Z = (Z2 - Z1) * c
+        Z = Z2 * c1 + Z1 * c2
 
         self.set_parameters(params)
         Z0 = op.heisenberg_tr(dev_wires, inverse=True)