docs/tutorials/gradients.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "xLOXFOT5Q40E"
   },
   "source": [
    "##### Copyright 2020 The TensorFlow Authors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {},
    "colab_type": "code",
    "id": "iiQkM5ZgQ8r2"
   },
   "outputs": [],
   "source": [
    "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#\n",
    "# https://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "j6331ZSsQGY3"
   },
   "source": [
    "# Calculate gradients"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "i9Jcnb8bQQyd"
   },
   "source": [
    "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
    "  <td>\n",
    "    <a target=\"_blank\" href=\"https://www.tensorflow.org/quantum/tutorials/gradients\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />View on TensorFlow.org</a>\n",
    "  </td>\n",
    "  <td>\n",
    "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/quantum/blob/master/docs/tutorials/gradients.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
    "  </td>\n",
    "  <td>\n",
    "    <a target=\"_blank\" href=\"https://github.com/tensorflow/quantum/blob/master/docs/tutorials/gradients.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
    "  </td>\n",
    "  <td>\n",
    "    <a href=\"https://storage.googleapis.com/tensorflow_docs/quantum/docs/tutorials/gradients.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\" />Download notebook</a>\n",
    "  </td>\n",
    "</table>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "FxkQA6oblNqI"
   },
   "source": [
    "This tutorial explores gradient calculation algorithms for the expectation values of quantum circuits.\n",
    "\n",
    "Calculating the gradient of the expectation value of a certain observable in a quantum circuit is an involved process. Expectation values of observables do not have the luxury of having analytic gradient formulas that are always easy to write down—unlike traditional machine learning transformations such as matrix multiplication or vector addition that have analytic gradient formulas which are easy to write down. As a result, there are different quantum gradient calculation methods that come in handy for different scenarios. This tutorial compares and contrasts two different differentiation schemes."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "pvG0gAJqGYJo"
   },
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "TorxE5tnkvb2"
   },
   "outputs": [],
   "source": [
    "!pip install tensorflow==2.15.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "OIbP5hklC338"
   },
   "source": [
    "Install TensorFlow Quantum:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "saFHsRDpkvkH"
   },
   "outputs": [],
   "source": [
    "!pip install tensorflow-quantum==0.7.3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "4Ql5PW-ACO0J"
   },
   "outputs": [],
   "source": [
    "# Update package resources to account for version changes.\n",
    "import importlib, pkg_resources\n",
    "\n",
    "importlib.reload(pkg_resources)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "MkTqyoSxGUfB"
   },
   "source": [
    "Now import TensorFlow and the module dependencies:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "enZ300Bflq80"
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import tensorflow_quantum as tfq\n",
    "\n",
    "import cirq\n",
    "import sympy\n",
    "import numpy as np\n",
    "\n",
    "# visualization tools\n",
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "from cirq.contrib.svg import SVGCircuit"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "b08Mmbs8lr81"
   },
   "source": [
    "## 1. Preliminary\n",
    "\n",
    "Let's make the notion of gradient calculation for quantum circuits a little more concrete. Suppose you have a parameterized circuit like this one:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "YkPYJ_Ak-GKu"
   },
   "outputs": [],
   "source": [
    "qubit = cirq.GridQubit(0, 0)\n",
    "my_circuit = cirq.Circuit(cirq.Y(qubit)**sympy.Symbol('alpha'))\n",
    "SVGCircuit(my_circuit)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "wgQIlCWy-MVr"
   },
   "source": [
    "Along with an observable:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "xurmJdFy-Jae"
   },
   "outputs": [],
   "source": [
    "pauli_x = cirq.X(qubit)\n",
    "pauli_x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "j3OzKYe5NT_W"
   },
   "source": [
    "Looking at this operator you know that $⟨Y(\\alpha)| X | Y(\\alpha)⟩ = \\sin(\\pi \\alpha)$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Ps-pd2mndXs7"
   },
   "outputs": [],
   "source": [
    "def my_expectation(op, alpha):\n",
    "    \"\"\"Compute ⟨Y(alpha)| `op` | Y(alpha)⟩\"\"\"\n",
    "    params = {'alpha': alpha}\n",
    "    sim = cirq.Simulator()\n",
    "    final_state_vector = sim.simulate(my_circuit, params).final_state_vector\n",
    "    return op.expectation_from_state_vector(final_state_vector, {qubit: 0}).real\n",
    "\n",
    "\n",
    "my_alpha = 0.3\n",
    "print(\"Expectation=\", my_expectation(pauli_x, my_alpha))\n",
    "print(\"Sin Formula=\", np.sin(np.pi * my_alpha))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "zcCX109cJUaz"
   },
   "source": [
    " and if you define $f_{1}(\\alpha) = ⟨Y(\\alpha)| X | Y(\\alpha)⟩$ then $f_{1}^{'}(\\alpha) = \\pi \\cos(\\pi \\alpha)$. Let's check this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "VMq7EayNRyQb"
   },
   "outputs": [],
   "source": [
    "def my_grad(obs, alpha, eps=0.01):\n",
    "    grad = 0\n",
    "    f_x = my_expectation(obs, alpha)\n",
    "    f_x_prime = my_expectation(obs, alpha + eps)\n",
    "    return ((f_x_prime - f_x) / eps).real\n",
    "\n",
    "\n",
    "print('Finite difference:', my_grad(pauli_x, my_alpha))\n",
    "print('Cosine formula:   ', np.pi * np.cos(np.pi * my_alpha))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "-SUlLpXBeicF"
   },
   "source": [
    "## 2. The need for a differentiator\n",
    "\n",
    "With larger circuits, you won't always be so lucky to have a formula that precisely calculates the gradients of a given quantum circuit. In the event that a simple formula isn't enough to calculate the gradient, the `tfq.differentiators.Differentiator` class allows you to define algorithms for computing the gradients of your circuits. For instance you can recreate the above example in TensorFlow Quantum (TFQ) with:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Om76ZLu8NT_i"
   },
   "outputs": [],
   "source": [
    "expectation_calculation = tfq.layers.Expectation(\n",
    "    differentiator=tfq.differentiators.ForwardDifference(grid_spacing=0.01))\n",
    "\n",
    "expectation_calculation(my_circuit,\n",
    "                        operators=pauli_x,\n",
    "                        symbol_names=['alpha'],\n",
    "                        symbol_values=[[my_alpha]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "lx3y2DX9NT_k"
   },
   "source": [
    "However, if you switch to estimating expectation based on sampling (what would happen on a true device) the values can change a little bit. This means you now have an imperfect estimate:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "v27rRyAHNT_l"
   },
   "outputs": [],
   "source": [
    "sampled_expectation_calculation = tfq.layers.SampledExpectation(\n",
    "    differentiator=tfq.differentiators.ForwardDifference(grid_spacing=0.01))\n",
    "\n",
    "sampled_expectation_calculation(my_circuit,\n",
    "                                operators=pauli_x,\n",
    "                                repetitions=500,\n",
    "                                symbol_names=['alpha'],\n",
    "                                symbol_values=[[my_alpha]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Igwa3EnzNT_p"
   },
   "source": [
    "This can quickly compound into a serious accuracy problem when it comes to gradients:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "StljXH38NT_q"
   },
   "outputs": [],
   "source": [
    "# Make input_points = [batch_size, 1] array.\n",
    "input_points = np.linspace(0, 5, 200)[:, np.newaxis].astype(np.float32)\n",
    "exact_outputs = expectation_calculation(my_circuit,\n",
    "                                        operators=pauli_x,\n",
    "                                        symbol_names=['alpha'],\n",
    "                                        symbol_values=input_points)\n",
    "imperfect_outputs = sampled_expectation_calculation(my_circuit,\n",
    "                                                    operators=pauli_x,\n",
    "                                                    repetitions=500,\n",
    "                                                    symbol_names=['alpha'],\n",
    "                                                    symbol_values=input_points)\n",
    "plt.title('Forward Pass Values')\n",
    "plt.xlabel('$x$')\n",
    "plt.ylabel('$f(x)$')\n",
    "plt.plot(input_points, exact_outputs, label='Analytic')\n",
    "plt.plot(input_points, imperfect_outputs, label='Sampled')\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "dfXObk7KNT_t"
   },
   "outputs": [],
   "source": [
    "# Gradients are a much different story.\n",
    "values_tensor = tf.convert_to_tensor(input_points)\n",
    "\n",
    "with tf.GradientTape() as g:\n",
    "    g.watch(values_tensor)\n",
    "    exact_outputs = expectation_calculation(my_circuit,\n",
    "                                            operators=pauli_x,\n",
    "                                            symbol_names=['alpha'],\n",
    "                                            symbol_values=values_tensor)\n",
    "analytic_finite_diff_gradients = g.gradient(exact_outputs, values_tensor)\n",
    "\n",
    "with tf.GradientTape() as g:\n",
    "    g.watch(values_tensor)\n",
    "    imperfect_outputs = sampled_expectation_calculation(\n",
    "        my_circuit,\n",
    "        operators=pauli_x,\n",
    "        repetitions=500,\n",
    "        symbol_names=['alpha'],\n",
    "        symbol_values=values_tensor)\n",
    "sampled_finite_diff_gradients = g.gradient(imperfect_outputs, values_tensor)\n",
    "\n",
    "plt.title('Gradient Values')\n",
    "plt.xlabel('$x$')\n",
    "plt.ylabel('$f^{\\'}(x)$')\n",
    "plt.plot(input_points, analytic_finite_diff_gradients, label='Analytic')\n",
    "plt.plot(input_points, sampled_finite_diff_gradients, label='Sampled')\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Ld34TJvTNT_w"
   },
   "source": [
    "Here you can see that although the finite difference formula is fast to compute the gradients themselves in the analytical case, when it came to the sampling based methods it was far too noisy. More careful techniques must be used to ensure a good gradient can be calculated. Next you will look at a much slower technique that wouldn't be as well suited for analytical expectation gradient calculations, but does perform much better in the real-world sample based case:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "JsBxH_RaNT_x"
   },
   "outputs": [],
   "source": [
    "# A smarter differentiation scheme.\n",
    "gradient_safe_sampled_expectation = tfq.layers.SampledExpectation(\n",
    "    differentiator=tfq.differentiators.ParameterShift())\n",
    "\n",
    "with tf.GradientTape() as g:\n",
    "    g.watch(values_tensor)\n",
    "    imperfect_outputs = gradient_safe_sampled_expectation(\n",
    "        my_circuit,\n",
    "        operators=pauli_x,\n",
    "        repetitions=500,\n",
    "        symbol_names=['alpha'],\n",
    "        symbol_values=values_tensor)\n",
    "\n",
    "sampled_param_shift_gradients = g.gradient(imperfect_outputs, values_tensor)\n",
    "\n",
    "plt.title('Gradient Values')\n",
    "plt.xlabel('$x$')\n",
    "plt.ylabel('$f^{\\'}(x)$')\n",
    "plt.plot(input_points, analytic_finite_diff_gradients, label='Analytic')\n",
    "plt.plot(input_points, sampled_param_shift_gradients, label='Sampled')\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "0xlUlh8wNT_z"
   },
   "source": [
    "From the above you can see that certain differentiators are best used for particular research scenarios. In general, the slower sample-based methods that are robust to device noise, etc., are great differentiators when testing or implementing algorithms in a more \"real world\" setting. Faster methods like finite difference are great for analytical calculations and you want higher throughput, but aren't yet concerned with the device viability of your algorithm."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "FaijzZ4MNT_0"
   },
   "source": [
    "## 3. Multiple observables\n",
    "\n",
    "Let's introduce a second observable and see how TensorFlow Quantum supports multiple observables for a single circuit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ytgB_DqDNT_3"
   },
   "outputs": [],
   "source": [
    "pauli_z = cirq.Z(qubit)\n",
    "pauli_z"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "r51TZls4NT_6"
   },
   "source": [
    "If this observable is used with the same circuit as before, then you have $f_{2}(\\alpha) = ⟨Y(\\alpha)| Z | Y(\\alpha)⟩ = \\cos(\\pi \\alpha)$ and $f_{2}^{'}(\\alpha) = -\\pi \\sin(\\pi \\alpha)$. Perform a quick check:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "19FKgu0ANT_7"
   },
   "outputs": [],
   "source": [
    "test_value = 0.\n",
    "\n",
    "print('Finite difference:', my_grad(pauli_z, test_value))\n",
    "print('Sin formula:      ', -np.pi * np.sin(np.pi * test_value))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "_33Y5mL0NT_-"
   },
   "source": [
    "It's a match (close enough).\n",
    "\n",
    "Now if you define $g(\\alpha) = f_{1}(\\alpha) + f_{2}(\\alpha)$ then $g'(\\alpha) = f_{1}^{'}(\\alpha) + f^{'}_{2}(\\alpha)$. Defining more than one observable in TensorFlow Quantum to use along with a circuit is equivalent to adding on more terms to $g$.\n",
    "\n",
    "This means that the gradient of a particular symbol in a circuit is equal to the sum of the gradients with regards to each observable for that symbol applied to that circuit. This is compatible with TensorFlow gradient taking and backpropagation (where you give the sum of the gradients over all observables as the gradient for a particular symbol)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "3WFJfFEbNT_-"
   },
   "outputs": [],
   "source": [
    "sum_of_outputs = tfq.layers.Expectation(\n",
    "    differentiator=tfq.differentiators.ForwardDifference(grid_spacing=0.01))\n",
    "\n",
    "sum_of_outputs(my_circuit,\n",
    "               operators=[pauli_x, pauli_z],\n",
    "               symbol_names=['alpha'],\n",
    "               symbol_values=[[test_value]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "-ujQUu3WNUAB"
   },
   "source": [
    "Here you see the first entry is the expectation w.r.t Pauli X, and the second is the expectation w.r.t Pauli Z. Now when you take the gradient:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "jcAQa9l0NUAB"
   },
   "outputs": [],
   "source": [
    "test_value_tensor = tf.convert_to_tensor([[test_value]])\n",
    "\n",
    "with tf.GradientTape() as g:\n",
    "    g.watch(test_value_tensor)\n",
    "    outputs = sum_of_outputs(my_circuit,\n",
    "                             operators=[pauli_x, pauli_z],\n",
    "                             symbol_names=['alpha'],\n",
    "                             symbol_values=test_value_tensor)\n",
    "\n",
    "sum_of_gradients = g.gradient(outputs, test_value_tensor)\n",
    "\n",
    "print(my_grad(pauli_x, test_value) + my_grad(pauli_z, test_value))\n",
    "print(sum_of_gradients.numpy())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "-fZmbYGANUAE"
   },
   "source": [
    "Here you have verified that the sum of the gradients for each observable is indeed the gradient of $\\alpha$. This behavior is supported by all TensorFlow Quantum differentiators and plays a crucial role in the compatibility with the rest of TensorFlow."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "lZsGG7lWNUAF"
   },
   "source": [
    "## 4. Advanced usage\n",
    "All differentiators that exist inside of TensorFlow Quantum subclass `tfq.differentiators.Differentiator`. To implement a differentiator, a user must implement one of two interfaces.  The standard is to implement `get_gradient_circuits`, which tells the base class which circuits to measure to obtain an estimate of the gradient.  Alternatively, you can overload `differentiate_analytic` and `differentiate_sampled`; the class `tfq.differentiators.Adjoint` takes this route.\n",
    "\n",
    "The following uses TensorFlow Quantum to implement the gradient of a circuit.  You will use a small example of parameter shifting."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "J1xN6Ln5mB9N"
   },
   "source": [
    "Recall the circuit you defined above, $|\\alpha⟩ = Y^{\\alpha}|0⟩$.  As before, you can define a function as the expectation value of this circuit against the $X$ observable, $f(\\alpha) = ⟨\\alpha|X|\\alpha⟩$.  Using [parameter shift rules](https://pennylane.ai/qml/glossary/parameter_shift.html), for this circuit, you can find that the derivative is\n",
    "$$\\frac{\\partial}{\\partial \\alpha} f(\\alpha) = \\frac{\\pi}{2} f\\left(\\alpha + \\frac{1}{2}\\right) -  \\frac{ \\pi}{2} f\\left(\\alpha - \\frac{1}{2}\\right)$$\n",
    "The `get_gradient_circuits` function returns the components of this derivative."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "5iY4q6FKNUAG"
   },
   "outputs": [],
   "source": [
    "class MyDifferentiator(tfq.differentiators.Differentiator):\n",
    "    \"\"\"A Toy differentiator for <Y^alpha | X |Y^alpha>.\"\"\"\n",
    "\n",
    "    def __init__(self):\n",
    "        pass\n",
    "\n",
    "    def get_gradient_circuits(self, programs, symbol_names, symbol_values):\n",
    "        \"\"\"Return circuits to compute gradients for given forward pass circuits.\n",
    "        \n",
    "        Every gradient on a quantum computer can be computed via measurements\n",
    "        of transformed quantum circuits.  Here, you implement a custom gradient\n",
    "        for a specific circuit.  For a real differentiator, you will need to\n",
    "        implement this function in a more general way.  See the differentiator\n",
    "        implementations in the TFQ library for examples.\n",
    "        \"\"\"\n",
    "\n",
    "        # The two terms in the derivative are the same circuit...\n",
    "        batch_programs = tf.stack([programs, programs], axis=1)\n",
    "\n",
    "        # ... with shifted parameter values.\n",
    "        shift = tf.constant(1 / 2)\n",
    "        forward = symbol_values + shift\n",
    "        backward = symbol_values - shift\n",
    "        batch_symbol_values = tf.stack([forward, backward], axis=1)\n",
    "\n",
    "        # Weights are the coefficients of the terms in the derivative.\n",
    "        num_program_copies = tf.shape(batch_programs)[0]\n",
    "        batch_weights = tf.tile(tf.constant([[[np.pi / 2, -np.pi / 2]]]),\n",
    "                                [num_program_copies, 1, 1])\n",
    "\n",
    "        # The index map simply says which weights go with which circuits.\n",
    "        batch_mapper = tf.tile(tf.constant([[[0, 1]]]),\n",
    "                               [num_program_copies, 1, 1])\n",
    "\n",
    "        return (batch_programs, symbol_names, batch_symbol_values,\n",
    "                batch_weights, batch_mapper)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "bvEgw2m6NUAI"
   },
   "source": [
    "The `Differentiator` base class uses the components returned from `get_gradient_circuits` to calculate the derivative, as in the parameter shift formula you saw above.  This new differentiator can now be used with existing `tfq.layer` objects:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "QrKnkWswNUAJ"
   },
   "outputs": [],
   "source": [
    "custom_dif = MyDifferentiator()\n",
    "custom_grad_expectation = tfq.layers.Expectation(differentiator=custom_dif)\n",
    "\n",
    "# Now let's get the gradients with finite diff.\n",
    "with tf.GradientTape() as g:\n",
    "    g.watch(values_tensor)\n",
    "    exact_outputs = expectation_calculation(my_circuit,\n",
    "                                            operators=[pauli_x],\n",
    "                                            symbol_names=['alpha'],\n",
    "                                            symbol_values=values_tensor)\n",
    "\n",
    "analytic_finite_diff_gradients = g.gradient(exact_outputs, values_tensor)\n",
    "\n",
    "# Now let's get the gradients with custom diff.\n",
    "with tf.GradientTape() as g:\n",
    "    g.watch(values_tensor)\n",
    "    my_outputs = custom_grad_expectation(my_circuit,\n",
    "                                         operators=[pauli_x],\n",
    "                                         symbol_names=['alpha'],\n",
    "                                         symbol_values=values_tensor)\n",
    "\n",
    "my_gradients = g.gradient(my_outputs, values_tensor)\n",
    "\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.title('Exact Gradient')\n",
    "plt.plot(input_points, analytic_finite_diff_gradients.numpy())\n",
    "plt.xlabel('x')\n",
    "plt.ylabel('f(x)')\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.title('My Gradient')\n",
    "plt.plot(input_points, my_gradients.numpy())\n",
    "plt.xlabel('x')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "oXqcJWigNUAL"
   },
   "source": [
    "This new differentiator can now be used to generate differentiable ops.\n",
    "\n",
    "Key Point: A differentiator that has been previously attached to an op must be refreshed before attaching to a new op, because a differentiator may only be attached to one op at a time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "F_WHcj3bNUAM"
   },
   "outputs": [],
   "source": [
    "# Create a noisy sample based expectation op.\n",
    "expectation_sampled = tfq.get_sampled_expectation_op(\n",
    "    cirq.DensityMatrixSimulator(noise=cirq.depolarize(0.01)))\n",
    "\n",
    "# Make it differentiable with your differentiator:\n",
    "# Remember to refresh the differentiator before attaching the new op\n",
    "custom_dif.refresh()\n",
    "differentiable_op = custom_dif.generate_differentiable_op(\n",
    "    sampled_op=expectation_sampled)\n",
    "\n",
    "# Prep op inputs.\n",
    "circuit_tensor = tfq.convert_to_tensor([my_circuit])\n",
    "op_tensor = tfq.convert_to_tensor([[pauli_x]])\n",
    "single_value = tf.convert_to_tensor([[my_alpha]])\n",
    "num_samples_tensor = tf.convert_to_tensor([[5000]])\n",
    "\n",
    "with tf.GradientTape() as g:\n",
    "    g.watch(single_value)\n",
    "    forward_output = differentiable_op(circuit_tensor, ['alpha'], single_value,\n",
    "                                       op_tensor, num_samples_tensor)\n",
    "\n",
    "my_gradients = g.gradient(forward_output, single_value)\n",
    "\n",
    "print('---TFQ---')\n",
    "print('Foward:  ', forward_output.numpy())\n",
    "print('Gradient:', my_gradients.numpy())\n",
    "print('---Original---')\n",
    "print('Forward: ', my_expectation(pauli_x, my_alpha))\n",
    "print('Gradient:', my_grad(pauli_x, my_alpha))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "OGWcpqzDNUAP"
   },
   "source": [
    "Success: Now you can use all the differentiators that TensorFlow Quantum has to offer—and define your own."
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "gradients.ipynb",
   "private_outputs": true,
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.10.9 (main, Dec  7 2022, 13:47:07) [GCC 12.2.0]"
  },
  "vscode": {
   "interpreter": {
    "hash": "916dbcbb3f70747c44a77c7bcd40155683ae19c65e1c03b4aa3499c5328201f1"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}