3-quantum-gradients.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "39hfro50I-ku",
    "outputId": "527446ae-3fd6-44d3-8a39-0224bac1ca4a"
   },
   "outputs": [],
   "source": [
    "%%capture \n",
    "\n",
    "# Comment this out if you don't want to install pennylane from this notebook\n",
    "!pip install pennylane"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "W9J7JpNNknG7"
   },
   "outputs": [],
   "source": [
    "import pennylane as qml\n",
    "from pennylane import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6i4quOW-I-kw"
   },
   "source": [
    "\n",
    "# Computing quantum gradients on any device\n",
    "\n",
    "In the last two notebooks we saw that PennyLane allows you to seamlessly \"slot\" quantum computations into automatic differentiation frameworks. While we used the Numpy/Autograd framework, you can also interface PennyLane with PyTorch, Tensorflow and Jax.\n",
    "\n",
    "If quantum computations are *simulated* by coding up matrix-vector multiplications in a specific coding language,  they just become differentiable computations. This is the trick that the `default.qubit` device uses if possible: it knows how to simulate quantum circuits in Autograd, TensorFlow etc. But what if we use a device that is not implemented in an automatic differentiation fashion? And what if the device uses quantum hardware?\n",
    "\n",
    "In PennyLane, any quantum device, whether a hardware device or a simulator, can be trained using the [parameter-shift rule](https://pennylane.ai/qml/glossary/parameter_shift.html) to compute quantum gradients. Indeed, the parameter-shift rule is ideally suited to hardware devices, as it does\n",
    "not require any knowledge about the internal workings of the device; it is sufficient to treat the device as a 'black box', and to query it with different input values in order to determine the gradient.\n",
    "\n",
    "## 1. The parameter-shift rule\n",
    "\n",
    "\n",
    "The parameter-shift rule states that, given a variational quantum circuit $U(\\boldsymbol\n",
    "\\theta)$ composed of parametrized Pauli rotations, and some measured observable $\\mathcal{M}$, the\n",
    "derivative of the expectation value\n",
    "\n",
    "\n",
    "\\begin{align}\\langle \\mathcal{M} \\rangle (\\boldsymbol\\theta) =\n",
    "    \\langle 0 \\vert U(\\boldsymbol\\theta)^\\dagger \\mathcal{M} U(\\boldsymbol\\theta) \\vert 0\\rangle\\end{align}\n",
    "\n",
    "with respect to the input circuit parameters $\\boldsymbol{\\theta}$ is given by\n",
    "\n",
    "\\begin{align}\\nabla_{\\theta_i}\\langle \\mathcal{M} \\rangle(\\boldsymbol\\theta)\n",
    "      =  \\frac{1}{2}\n",
    "            \\left[\n",
    "                \\langle \\mathcal{M} \\rangle\\left(\\boldsymbol\\theta + \\frac{\\pi}{2}\\hat{\\mathbf{e}}_i\\right)\n",
    "              - \\langle \\mathcal{M} \\rangle\\left(\\boldsymbol\\theta - \\frac{\\pi}{2}\\hat{\\mathbf{e}}_i\\right)\n",
    "            \\right].\\end{align}\n",
    "\n",
    "Thus, the gradient of the expectation value can be calculated by evaluating the same variational quantum circuit, but with shifted parameter values (hence the name, parameter-shift rule!).\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "iFOUl-CnI-ky"
   },
   "source": [
    "## 2. Hand-coded parameter-shift rule"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jcFzzMPhI-k1"
   },
   "source": [
    "Let's first have a go implementing the parameter-shift rule manually. We use the quantum model from the previous notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "id": "FqTDlFNfI-kz"
   },
   "outputs": [],
   "source": [
    "dev = qml.device(\"default.qubit\", wires=2) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "q-OLTiSII-k1",
    "outputId": "ab719876-147d-4406-819f-49b6b151a876"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0.012023, requires_grad=True)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "@qml.qnode(dev, diff_method='parameter-shift')\n",
    "def quantum_model(x, w):\n",
    "    qml.templates.AngleEmbedding(x, wires=[0, 1])\n",
    "    qml.templates.BasicEntanglerLayers(w, wires=[0, 1])\n",
    "    return qml.expval(qml.PauliZ(wires=0))\n",
    "\n",
    "x = np.array([0.1, 0.2], requires_grad=False)\n",
    "w = np.array([[-2.1, 1.2], [-1.4, -3.9], [0.5, 0.2]])\n",
    "\n",
    "quantum_model(x, w)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "WBDp8Vg9spMv",
    "outputId": "7c2e0795-0900-4315-db15-b77c36a8b30b"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 0: ──RX(0.1)──RX(-2.1)──╭C──RX(-1.4)──╭C──RX(0.5)──╭C──┤ ⟨Z⟩ \n",
      " 1: ──RX(0.2)──RX(1.2)───╰X──RX(-3.9)──╰X──RX(0.2)──╰X──┤     \n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(quantum_model.draw())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "J491XZp6I-k3"
   },
   "source": [
    "Now that we have defined our variational circuit QNode, we can construct\n",
    "a function that computes the gradient of the $i\\text{th}$ parameter\n",
    "using the parameter-shift rule. Since the model also depends on input $x$, the gradient will be computed for this fixed $x$ only.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "BMIoXqxdI-k4",
    "outputId": "ca683e7a-be2d-42c9-c5ae-8ed30b4de6eb"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.16954137293040533\n"
     ]
    }
   ],
   "source": [
    "def parameter_shift_term(qnode, params, x, i, j):\n",
    "    \n",
    "    shifted = params.copy()\n",
    "    \n",
    "    shifted[i, j] += np.pi/2\n",
    "    forward = qnode(x, shifted)  # forward evaluation\n",
    "\n",
    "    shifted[i, j] -= np.pi\n",
    "    backward = qnode(x, shifted) # backward evaluation\n",
    "\n",
    "    return 0.5 * (forward - backward)\n",
    "\n",
    "# gradient with respect to the first parameter\n",
    "print(parameter_shift_term(quantum_model, w, x, 0, 0))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZzZ5um5VI-k4"
   },
   "source": [
    "In order to compute the gradient with respect to *all* parameters, we need\n",
    "to loop over the indices ``i`` and ``j``:\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "1k_45vk_I-k4",
    "outputId": "cccb11ec-0d17-4af6-f65b-af4209653b5d"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[1.69541373e-01 0.00000000e+00]\n",
      " [6.97079563e-02 1.11022302e-16]\n",
      " [1.69541373e-01 0.00000000e+00]]\n"
     ]
    }
   ],
   "source": [
    "def parameter_shift(qnode, params, x):\n",
    "    gradients = np.zeros_like((params))\n",
    "\n",
    "    for i in range(len(gradients)):\n",
    "        for j in range(len(gradients[0])):\n",
    "            gradients[i, j] = parameter_shift_term(qnode, w, x, i, j)\n",
    "\n",
    "    return gradients\n",
    "\n",
    "print(parameter_shift(quantum_model, w, x))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Built-in parameter shift differentiation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "RToLvQIFI-k5"
   },
   "source": [
    "We can compare this to PennyLane's *built-in* parameter-shift feature by using\n",
    "the `qml.grad` function. Remember, when we defined the\n",
    "QNode, we specified that we wanted it to be differentiable using the parameter-shift\n",
    "method (``diff_method=\"parameter-shift\"``).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "v21125ecI-k6",
    "outputId": "3e0ad358-7780-481d-80fd-f7d585990292"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[1.69541373e-01 0.00000000e+00]\n",
      " [6.97079563e-02 1.11022302e-16]\n",
      " [1.69541373e-01 0.00000000e+00]]\n"
     ]
    }
   ],
   "source": [
    "grad_function = qml.grad(quantum_model, argnum=1)\n",
    "\n",
    "print(grad_function(x, w))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5OZsKssBI-k6"
   },
   "source": [
    "If you count the number of quantum evaluations, you will notice that we had to evaluate the circuit roughly\n",
    "``2*len(params)`` number of times in order to compute the quantum gradient with respect to all\n",
    "parameters. \n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "RdHKpihpvCYZ",
    "outputId": "2abedb90-e099-4a0c-8dac-729ea06b4922"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of evaluations per gradient: 13\n"
     ]
    }
   ],
   "source": [
    "before = dev.num_executions \n",
    "grad_function = qml.grad(quantum_model, argnum=1)\n",
    "grad_function(x, w)\n",
    "after = dev.num_executions\n",
    "print(\"Number of evaluations per gradient:\", after-before)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "LpO7WcDjvBV6"
   },
   "source": [
    "While reasonably fast for a small number of parameters, as the number of parameters in\n",
    "our quantum circuit grows, so does both\n",
    "\n",
    "1. the circuit depth (and thus the time taken to evaluate each expectation value or 'forward' pass), and\n",
    "\n",
    "2. the number of parameter-shift evaluations required.\n",
    "\n",
    "Both of these factors increase the time taken to compute the gradient with\n",
    "respect to all parameters.\n",
    "\n",
    "It is therefore crucial to develop efficient pipelines for the evaluation of gradients in quantum machine learning. If you want to learn more about this, check out the [Amazon-Braket demo](https://pennylane.ai/qml/demos/braket-parallel-gradients.html), which explains how PennyLane and AWS have teamed up to paralellise the evaluation of quantum gradients."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "UtK1fbrVI-k7"
   },
   "source": [
    "#  TASKS\n",
    "\n",
    "1. Create a new `default.qubit` device where you can set the `diff_method` explicitely to `diff_method = 'backprop'`. This will use automatic differentiation tricks to compute gradients much more efficiently for simulations (as implicitely used in the previous notebooks). Find out how often the device is evaluated when you compute a gradient using the `qml.grad` function. \n",
    "\n",
    "2. In each step of gradient descent we need to compute the gradient of a quantum computation. Write a function `num_evals(n_params, n_steps)` that takes the number of parameters as well as the number of steps and returns the number of circuit evaluations needed for gradient descent training with a parameter shift rule. Let's say you need as many training steps as you have parameters or `n_steps=n_params`. Plot the number of circuit evaluations over the number of parameters. What scaling behaviour do you see? Compare this to a scaling in which each step only takes a single evaluation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ca6JthZCwpKj"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "name": "3-quantum-gradients.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "cern_tutorial",
   "language": "python",
   "name": "cern_tutorial"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}