text: # Copyright 2018-2020 Xanadu Quantum Technologies Inc.
  
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  
  #     http://www.apache.org/licenses/LICENSE-2.0
  
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  """
  This module contains the ``JacobianTape``, which adds differentiation methods
  to the ``QuantumTape`` class.
  """
  # pylint: disable=too-many-branches
  
  import itertools
  import numpy as np
  
  import pennylane as qml
  
  from pennylane.operation import State
  from pennylane.tape.tapes.tape import QuantumTape
  
  # CV ops still need to support state preparation operations prior to any
  # other operation for PennyLane-SF tests to pass.
  STATE_PREP_OPS = (
      qml.BasisState,
      qml.QubitStateVector,
      # qml.CatState,
      # qml.CoherentState,
      # qml.FockDensityMatrix,
      # qml.DisplacedSqueezedState,
      # qml.FockState,
      # qml.FockStateVector,
      # qml.ThermalState,
      # qml.GaussianState,
  )
  
  
  # pylint: disable=too-many-public-methods
  class JacobianTape(QuantumTape):
      """A quantum tape recorder, that records, validates, executes,
      and differentiates variational quantum programs.
  
      .. note::
  
          As the quantum tape is a *beta* feature. See :mod:`pennylane.tape`
          for more details.
  
      Args:
          name (str): a name given to the quantum tape
  
      **Example**
  
      In addition to the functionality of the ``QuantumTape`` class,
      the ``JacobianTape`` has the ability to compute Jacobians of a quantum
      circuit. Consider the following
  
      .. code-block:: python
  
          import pennylane.tape
  
          with qml.tape.JacobianTape() as tape:
              qml.RX(0.432, wires=0)
              qml.RY(0.543, wires=0)
              qml.CNOT(wires=[0, 'a'])
              qml.RX(0.133, wires='a')
              qml.expval(qml.PauliZ(wires=[0]))
  
      The Jacobian is computed using finite difference:
  
      >>> tape.jacobian(dev)
      [[-0.35846484 -0.46923704  0.        ]]
      >>> tape.jacobian(dev, params=[0.1, 0.1, 0.1])
      [[-0.09933471 -0.09933471  0.        ]]
  
      Trainable parameters are taken into account when calculating the Jacobian,
      avoiding unnecessary calculations:
  
      >>> tape.trainable_params = {0} # set only the first parameter as free
      >>> tape.set_parameters(0.56)
      >>> tape.jacobian(dev)
      [[-0.45478169]]
      """
  
      def __init__(self, name=None):
          super().__init__(name=name)
          self.jacobian_options = {}
  
      def _grad_method(self, idx, use_graph=True, default_method="F"):
          """Determine the correct partial derivative computation method for each gate parameter.
  
          Parameter gradient methods include:
  
          * ``None``: the parameter does not support differentiation.
  
          * ``"0"``: the variational circuit output does not depend on this
            parameter (the partial derivative is zero).
  
          * ``"F"``: the parameter has a non-zero derivative that should be computed
            using finite-differences.
  
          * ``"A"``: the parameter has a non-zero derivative that should be computed
            using an analytic method.
  
          .. note::
  
              The base ``JacobianTape`` class only supports numerical differentiation, so
              this method will always return either ``"F"`` or ``None``. If an inheriting
              tape supports analytic differentiation for certain operations, make sure
              that this method is overwritten appropriately to return ``"A"`` where
              required.
  
          Args:
              idx (int): parameter index
              use_graph: whether to use a directed-acyclic graph to determine
                  if the parameter has a gradient of 0
              default_method (str): the default differentiation value to return
                  for parameters that where the grad method is not ``"0"`` or ``None``
  
          Returns:
              str: partial derivative method to be used
          """
          op = self._par_info[idx]["op"]
  
          if op.grad_method is None:
              return None
  
          if (self._graph is not None) or use_graph:
              # an empty list to store the 'best' partial derivative method
              # for each observable
              best = []
  
              # loop over all observables
              for ob in self.observables:
                  # check if op is an ancestor of ob
                  has_path = self.graph.has_path(op, ob)
  
                  # Use finite differences if there is a path, else the gradient is zero
                  best.append(default_method if has_path else "0")
  
              if all(k == "0" for k in best):
                  return "0"
  
          return default_method
  
      def _update_gradient_info(self):
          """Update the parameter information dictionary with gradient information
          of each parameter"""
  
          for i, info in self._par_info.items():
  
              if i not in self.trainable_params:
                  info["grad_method"] = None
              else:
                  info["grad_method"] = self._grad_method(i, use_graph=True)
  
      def _grad_method_validation(self, method):
          """Validates if the gradient method requested is supported by the trainable
          parameters, and returns the allowed parameter gradient methods.
  
          This method will generate parameter gradient information if it has not already
          been generated, and then proceed to validate the gradient method. In particular:
  
          * An exception will be raised if there exist non-differentiable trainable
            parameters on the tape.
  
          * An exception will be raised if the Jacobian method is ``"analytic"`` but there
            exist some trainable parameters on the tape that only support numeric differentiation.
  
          If all validations pass, this method will return a tuple containing the allowed parameter
          gradient methods for each trainable parameter.
  
          Args:
              method (str): the overall Jacobian differentiation method
  
          Returns:
              tuple[str, None]: the allowed parameter gradient methods for each trainable parameter
          """
  
          if "grad_method" not in self._par_info[0]:
              self._update_gradient_info()
  
          diff_methods = {
              idx: info["grad_method"]
              for idx, info in self._par_info.items()
              if idx in self.trainable_params
          }
  
          # check and raise an error if any parameters are non-differentiable
          nondiff_params = {idx for idx, g in diff_methods.items() if g is None}
  
          if nondiff_params:
              raise ValueError(f"Cannot differentiate with respect to parameter(s) {nondiff_params}")
  
          numeric_params = {idx for idx, g in diff_methods.items() if g == "F"}
  
          # If explicitly using analytic mode, ensure that all parameters
          # support analytic differentiation.
          if method == "analytic" and numeric_params:
              raise ValueError(
                  f"The analytic gradient method cannot be used with the argument(s) {numeric_params}."
              )
  
          return tuple(diff_methods.values())
  
      @staticmethod
      def _has_trainable_params(params, diff_methods):
          """Determines if there are any trainable parameters.
  
          Args:
              params (array[float]): one-dimensional array of parameters
              diff_methods (Sequence[str]): The corresponding differentiation method for each parameter.
                  A differentiation method of ``"0"`` corresponds to a constant parameter.
          """
          return params.size and not all(g == "0" for g in diff_methods)
  
      @staticmethod
      def _flatten_processing_result(g):
          """Flattens the output from processing_fn in parameter shift methods."""
          if g.dtype is np.dtype("object"):
              # object arrays cannot be flattened; must hstack them
              g = np.hstack(g)
  
          return g.flatten()
  
      def numeric_pd(self, idx, params=None, **options):
          """Generate the tapes and postprocessing methods required to compute the gradient of a parameter using the
          finite-difference method.
  
          Args:
              idx (int): trainable parameter index to differentiate with respect to
              params (list[Any]): The quantum tape operation parameters. If not provided,
                 the current tape parameter values are used (via :meth:`~.get_parameters`).
  
          Keyword Args:
              h=1e-7 (float): finite difference method step size
              order=1 (int): The order of the finite difference method to use. ``1`` corresponds
                  to forward finite differences, ``2`` to centered finite differences.
  
          Returns:
              tuple[list[QuantumTape], function]: A tuple containing the list of generated tapes,
              in addition to a post-processing function to be applied to the evaluated
              tapes.
          """
  
          if params is None:
              params = np.array(self.get_parameters())
  
          order = options.get("order", 1)
          h = options.get("h", 1e-7)
  
          shift = np.zeros_like(params, dtype=np.float64)
          shift[idx] = h
  
          if order == 1:
              # forward finite-difference.
  
              tapes = []
  
              # get the stored result of the original circuit
              y0 = options.get("y0", None)
  
              shifted = self.copy(copy_operations=True, tape_cls=QuantumTape)
              shifted.set_parameters(params + shift)
  
              tapes.append(shifted)
  
              if y0 is None:
                  tapes.append(self)
  
              def processing_fn(results):
                  """Computes the gradient of the parameter at index idx via first-order
                  forward finite differences.
  
                  Args:
                      results (list[real]): evaluated quantum tapes
  
                  Returns:
                      array[float]: 1-dimensional array of length determined by the tape output
                      measurement statistics
                  """
                  shifted = np.array(results[0])
                  unshifted = y0
  
                  if unshifted is None:
                      unshifted = np.array(results[1])
  
                  return (shifted - unshifted) / h
  
              return tapes, processing_fn
  
          if order == 2:
              # central finite difference
  
              shifted_forward = self.copy(copy_operations=True, tape_cls=QuantumTape)
              shifted_forward.set_parameters(params + shift / 2)
  
              shifted_backward = self.copy(copy_operations=True, tape_cls=QuantumTape)
              shifted_backward.set_parameters(params - shift / 2)
  
              tapes = [shifted_forward, shifted_backward]
  
              def second_order_processing_fn(results):
                  """Computes the gradient of the parameter at index idx via second-order
                  centered finite differences.
  
                  Args:
                      results (list[real]): evaluated quantum tapes
  
                  Returns:
                      array[float]: 1-dimensional array of length determined by the tape output
                      measurement statistics
                  """
                  res0 = np.array(results[0])
                  res1 = np.array(results[1])
                  return (res0 - res1) / h
  
              return tapes, second_order_processing_fn
  
          raise ValueError("Order must be 1 or 2.")
  
      def device_pd(self, device, params=None, **options):
          """Evaluate the gradient of the tape with respect to
          all trainable tape parameters by querying the provided device.
  
          Args:
              device (.Device, .QubitDevice): a PennyLane device
                  that can execute quantum operations and return measurement statistics
              params (list[Any]): The quantum tape operation parameters. If not provided,
                  the current tape parameter values are used (via :meth:`~.get_parameters`).
          """
          jacobian_method = getattr(device, options.get("jacobian_method", "jacobian"))
  
          if params is None:
              params = np.array(self.get_parameters())
  
          saved_parameters = self.get_parameters()
  
          # temporarily mutate the in-place parameters
          self.set_parameters(params)
  
          # TODO: modify devices that have device Jacobian methods to
          # accept the quantum tape as an argument
          jac = jacobian_method(self)
  
          # restore original parameters
          self.set_parameters(saved_parameters)
          return jac
  
      def analytic_pd(self, idx, params, **options):
          """Generate the quantum tapes and classical post-processing function required to compute the
          gradient of the tape with respect to a single trainable tape parameter using an analytic
          method.
  
          Args:
              idx (int): trainable parameter index to differentiate with respect to
              params (list[Any]): the quantum tape operation parameters
  
          Returns:
              tuple[list[QuantumTape], function]: A tuple containing the list of generated tapes,
              in addition to a post-processing function to be applied to the evaluated
              tapes.
          """
          raise NotImplementedError
  
      def hessian_pd(self, i, j, params, **options):
          """Generate the quantum tapes and classical post-processing function required to compute the
          Hessian of the tape with respect to two trainable tape parameter using an analytic
          method.
  
          Args:
              i (int): trainable parameter index to differentiate with respect to
              j (int): trainable parameter index to differentiate with respect to
              params (list[Any]): the quantum tape operation parameters
  
          Returns:
              tuple[list[QuantumTape], function]: A tuple containing the list of generated tapes,
              in addition to a post-processing function to be applied to the evaluated
              tapes.
          """
          raise NotImplementedError
  
      def jacobian(self, device, params=None, **options):
          r"""Compute the Jacobian of the parametrized quantum circuit recorded by the quantum tape.
  
          The quantum tape can be interpreted as a simple :math:`\mathbb{R}^m \to \mathbb{R}^n` function,
          mapping :math:`m` (trainable) gate parameters to :math:`n` measurement statistics,
          such as expectation values or probabilities.
  
          By default, the Jacobian will be computed with respect to all parameters on the quantum tape.
          This can be modified by setting the :attr:`~.trainable_params` attribute of the tape.
  
          The Jacobian can be computed using several methods:
  
          * Finite differences (``'numeric'``). The first-order method evaluates the circuit at
            :math:`n+1` points of the parameter space, the second-order method at :math:`2n` points,
            where ``n = tape.num_params``.
  
          * Analytic method (``'analytic'``). Analytic, if implemented by the inheriting quantum tape.
  
          * Best known method for each parameter (``'best'``): uses the analytic method if
            possible, otherwise finite difference.
  
          * Device method (``'device'``): Delegates the computation of the Jacobian to the
            device executing the circuit. Only supported by devices that provide their
            own method for computing derivatives; support can be checked by
            querying the device capabilities: ``dev.capabilities()['provides_jacobian']`` must
            return ``True``. Examples of supported devices include the experimental
            :class:`"default.tensor.tf" <~.DefaultTensorTF>` device.
  
          .. note::
  
              The finite difference method is sensitive to statistical noise in the circuit output,
              since it compares the output at two points infinitesimally close to each other. Hence
              the ``'F'`` method works best with exact expectation values when using simulator
              devices.
  
          Args:
              device (.Device, .QubitDevice): a PennyLane device
                  that can execute quantum operations and return measurement statistics
              params (list[Any]): The quantum tape operation parameters. If not provided,
                  the current tape parameter values are used (via :meth:`~.get_parameters`).
  
          Keyword Args:
              method="best" (str): The differentiation method. Must be one of ``"numeric"``,
                  ``"analytic"``, ``"best"``, or ``"device"``.
              h=1e-7 (float): finite difference method step size
              order=1 (int): The order of the finite difference method to use. ``1`` corresponds
                  to forward finite differences, ``2`` to centered finite differences.
              shift=pi/2 (float): the size of the shift for two-term parameter-shift gradient computations
  
          Returns:
              array[float]: 2-dimensional array of shape ``(tape.output_dim, tape.num_params)``
  
          **Example**
  
          .. code-block:: python
  
              with JacobianTape() as tape:
                  qml.RX(0.432, wires=0)
                  qml.RY(0.543, wires=0)
                  qml.CNOT(wires=[0, 'a'])
                  qml.RX(0.133, wires='a')
                  qml.probs(wires=[0, 'a'])
  
          If parameters are not provided, the existing tape parameters are used:
  
          >>> dev = qml.device("default.qubit", wires=[0, 'a'])
          >>> tape.jacobian(dev)
          array([[-0.178441  , -0.23358253, -0.05892804],
                 [-0.00079144, -0.00103601,  0.05892804],
                 [ 0.00079144,  0.00103601,  0.00737611],
                 [ 0.178441  ,  0.23358253, -0.00737611]])
  
          Parameters can be optionally passed during execution:
  
          >>> tape.jacobian(dev, params=[1.0, 0.0, 1.0])
          array([[-3.24029934e-01, -9.99200722e-09, -3.24029934e-01],
                 [-9.67055711e-02, -2.77555756e-09,  3.24029935e-01],
                 [ 9.67055709e-02,  3.05311332e-09,  9.67055709e-02],
                 [ 3.24029935e-01,  1.08246745e-08, -9.67055711e-02]])
  
          Parameters provided for execution are temporary, and do not affect
          the tapes' parameters in-place:
  
          >>> tape.get_parameters()
          [0.432, 0.543, 0.133]
  
          Explicitly setting the trainable parameters can significantly reduce
          computational resources, as non-trainable parameters are ignored
          during the computation:
  
          >>> tape.trainable_params = {0} # set only the first parameter as trainable
          >>> tape.jacobian(dev)
          array([[-0.178441  ],
                 [-0.00079144],
                 [ 0.00079144],
                 [ 0.178441  ]])
  
          If a tape has no trainable parameters, the Jacobian will be empty:
  
          >>> tape.trainable_params = {}
          >>> tape.jacobian(dev)
          array([], shape=(4, 0), dtype=float64)
          """
          if any([m.return_type is State for m in self.measurements]):
              raise ValueError("The jacobian method does not support circuits that return the state")
  
          if self.is_sampled:
              raise qml.QuantumFunctionError(
                  "Circuits that include sampling can not be differentiated."
              )
  
          method = options.get("method", "best")
  
          if method not in ("best", "numeric", "analytic", "device"):
              raise ValueError(f"Unknown gradient method '{method}'")
  
          if params is None:
              params = self.get_parameters()
  
          params = np.array(params)
  
          if method == "device":
              # Using device mode; simply query the device for the Jacobian
              return self.device_pd(device, params=params, **options)
  
          # perform gradient method validation
          diff_methods = self._grad_method_validation(method)
  
          if not self._has_trainable_params(params, diff_methods):
              # Either all parameters have grad method 0, or there are no trainable
              # parameters. Simply return an empty Jacobian.
              return np.zeros((self.output_dim, len(params)), dtype=float)
  
          if method == "numeric" or "F" in diff_methods:
              # there exist parameters that will be differentiated numerically
  
              if options.get("order", 1) == 1:
                  # First order (forward) finite-difference will be performed.
                  # Compute the value of the tape at the current parameters here. This ensures
                  # this computation is only performed once, for all parameters.
                  options["y0"] = np.asarray(self.execute_device(params, device))
  
          # some gradient methods need the device or the device wires
          options["device"] = device
          options["dev_wires"] = device.wires
  
          # collect all circuits (tapes) and postprocessing functions required
          # to compute the jacobian
  
          all_tapes = []
          reshape_info = []
          processing_fns = []
          nonzero_grad_idx = []
  
          for trainable_idx, param_method in enumerate(diff_methods):
              if param_method == "0":
                  continue
  
              nonzero_grad_idx.append(trainable_idx)
  
              if (method == "best" and param_method[0] == "F") or (method == "numeric"):
                  # numeric method
                  tapes, processing_fn = self.numeric_pd(trainable_idx, params=params, **options)
  
              elif (method == "best" and param_method[0] == "A") or (method == "analytic"):
                  # analytic method
                  tapes, processing_fn = self.analytic_pd(trainable_idx, params=params, **options)
  
              processing_fns.append(processing_fn)
  
              # we create a flat list here to feed at once to the device
              all_tapes.extend(tapes)
  
              # to extract the correct result for this parameter later, remember the number of tapes
              reshape_info.append(len(tapes))
  
          # execute all tapes at once
          results = device.batch_execute(all_tapes)
  
          # post-process the results with the appropriate function to fill jacobian columns with gradients
          jac = None
          start = 0
  
          for i, processing_fn, res_len in zip(nonzero_grad_idx, processing_fns, reshape_info):
              # extract the correct results from the flat list
              res = results[start : start + res_len]
              start += res_len
  
              # postprocess results to compute the gradient
              g = self._flatten_processing_result(processing_fn(res))
  
              if jac is None:
                  # update the tape's output dimension
                  self._output_dim = len(g)
                  # create the Jacobian matrix
                  jac = np.zeros((len(g), len(params)), dtype=float)
  
              jac[:, i] = g
  
          return jac
  
      def hessian(self, device, params=None, **options):
          r"""Compute the Hessian of the parametrized quantum circuit recorded by the quantum tape.
  
          The quantum tape can be interpreted as a simple :math:`\mathbb{R}^m \to \mathbb{R}^n` function,
          mapping :math:`m` (trainable) gate parameters to :math:`n` measurement statistics,
          such as expectation values or probabilities.
  
          By default, the Hessian will be computed with respect to all parameters on the quantum tape.
          This can be modified by setting the :attr:`~.trainable_params` attribute of the tape.
  
          The Hessian can be currently computed using only the ``'analytic'`` method.
  
          Args:
              device (.Device, .QubitDevice): a PennyLane device
                  that can execute quantum operations and return measurement statistics
              params (list[Any]): The quantum tape operation parameters. If not provided,
                  the current tape parameter values are used (via :meth:`~.get_parameters`).
  
          Keyword Args:
              method="analytic" (str): The differentiation method. Currently only
                  supports ``"analytic"``.
              s1=pi/2 (float): the size of the shift for index i in the parameter-shift Hessian computations
              s2=pi/2 (float): the size of the shift for index j in the parameter-shift Hessian computations
  
          Returns:
              array[float]: 2-dimensional array of shape ``(tape.num_params, tape.num_params)``
  
          **Example**
  
          .. code-block:: python
  
              n_wires = 5
              weights = [2.73943676, 0.16289932, 3.4536312, 2.73521126, 2.6412488]
  
              with QubitParamShiftTape() as tape:
                  for i in range(n_wires):
                      qml.RX(weights[i], wires=i)
  
                  qml.CNOT(wires=[0, 1])
                  qml.CNOT(wires=[2, 1])
                  qml.CNOT(wires=[3, 1])
                  qml.CNOT(wires=[4, 3])
  
                  qml.expval(qml.PauliZ(1))
  
          If parameters are not provided, the existing tape parameters are used:
  
          >>> dev = qml.device("default.qubit", wires=n_wires)
          >>> tape.hessian(dev)
          array([[ 0.79380556,  0.05549219,  0.10891309, -0.1452963,   0.],
                 [ 0.05549219,  0.79380556, -0.04208544,  0.05614438,  0.],
                 [ 0.10891309, -0.04208544,  0.79380556,  0.11019314,  0.],
                 [-0.1452963,   0.05614438,  0.11019314,  0.79380556,  0.],
                 [ 0.,          0.,          0.,          0.,          0.]])
  
          Parameters can be optionally passed during execution:
  
          >>> tape.hessian(dev, params=[1.0, 1.0, 2.0, 0, 0])
          array([[ 0.12148432, -0.29466251,  0.41341091,  0.,          0.],
                 [-0.29466251,  0.12148432,  0.41341091,  0.,          0.],
                 [ 0.41341091,  0.41341091,  0.12148432,  0.,          0.],
                 [ 0.,          0.,          0.,          0.12148432,  0.],
                 [ 0.,          0.,          0.,          0.,          0.]])
  
          Parameters provided for execution are temporary, and do not affect
          the tapes' parameters in-place:
  
          >>> tape.get_parameters()
          [2.73943676, 0.16289932, 3.4536312, 2.73521126, 2.6412488]
  
          If a tape has no trainable parameters, the Hessian will be empty:
  
          >>> tape.trainable_params = {}
          >>> tape.hessian(dev)
          array([], shape=(0, 0), dtype=float64)
          """
          if any([m.return_type is State for m in self.measurements]):
              raise ValueError("The Hessian method does not support circuits that return the state")
  
          method = options.get("method", "analytic")
  
          if method != "analytic":
              raise ValueError(f"Unknown Hessian method '{method}'")
  
          if params is None:
              params = self.get_parameters()
  
          params = np.array(params)
  
          # perform gradient method validation
          diff_methods = self._grad_method_validation(method)
  
          if not self._has_trainable_params(params, diff_methods):
              # Either all parameters have grad method 0, or there are no trainable
              # parameters. Simply return an empty Hessian.
              return np.zeros((len(params), len(params)), dtype=float)
  
          # some gradient methods need the device or the device wires
          options["device"] = device
          options["dev_wires"] = device.wires
  
          # collect all circuits (tapes) and postprocessing functions required
          # to compute the Hessian
          all_tapes = []
          reshape_info = []
          processing_fns = []
          nonzero_grad_idx = []
  
          # From Schwarz's theorem, the Hessian will be symmetric, so we
          # can compute the upper triangular part only and symmetrize
          # the final Hessian.
          for i, j in itertools.combinations_with_replacement(range(len(diff_methods)), 2):
              if diff_methods[i] == "0" or diff_methods[j] == "0":
                  continue
  
              nonzero_grad_idx.append((i, j))
  
              tapes, processing_fn = self.hessian_pd(i, j, params=params, **options)
  
              processing_fns.append(processing_fn)
  
              # we create a flat list here to feed at once to the device
              all_tapes.extend(tapes)
  
              # to extract the correct result for this parameter later, remember the number of tapes
              reshape_info.append(len(tapes))
  
          # execute all tapes at once
          results = device.batch_execute(all_tapes)
  
          hessian = None
          start = 0
  
          for (i, j), processing_fn, res_len in zip(nonzero_grad_idx, processing_fns, reshape_info):
              # extract the correct results from the flat list
              res = results[start : start + res_len]
              start += res_len
  
              # postprocess results to compute the gradient
              g = self._flatten_processing_result(processing_fn(res))
  
              if hessian is None:
                  # create the Hessian matrix
                  hessian = np.zeros((len(params), len(params)), dtype=float)
  
              if i == j:
                  hessian[i, i] = g
              else:
                  hessian[i, j] = hessian[j, i] = g
  
          return hessian
  