Adjoint_Differentiation.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "view-in-github"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/MonitSharma/Learn-Quantum-Machine-Learning/blob/main/Adjoint_Differentiation.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LEqORNrkqHBb"
      },
      "source": [
        "# Adjoint Differentiation\n",
        "\n",
        "[Classical automatic differentiation](https://en.wikipedia.org/wiki/Automatic_differentiation#The_chain_rule,_forward_and_reverse_accumulation) has two methods for calculation\n",
        "* forward\n",
        "* reverse\n",
        "\n",
        "\n",
        "\n",
        "The Optimal choice of method depends on the structutre of the problem, whether the function is **many-to-one** or **one-to-many**?\n",
        "\n",
        "\n",
        "\n",
        "Most methods for calculating the derivatives of quantum circuit are either direct applications of classical gradient methods to quantum simulations, or quantum hardware methods like parameter-shift where we can only extract restricted pieces of information.\n",
        "\n",
        "------\n",
        "\n",
        "Adjoint differentiation studies these two strategies, taking benefits from each. On simulators, we can examine and modify the state vector at any point. At the same time, we know our quantum circuit holds specific properties not present in the arbitrary classsical computation.\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "Quantum circuits only involve:\n",
        "\n",
        "1. initialization $$|0⟩$$\n",
        "\n",
        "2. application of unitary operators $$|\\Psi\\rangle = U_{n} U_{n-1} \\dots U_0 |0\\rangle,$$\n",
        "\n",
        "3. measurement, such as estimating an expectation value of a Hermitian operator $$\\langle M \\rangle = \\langle \\Psi | M | \\Psi \\rangle.$$\n",
        "\n",
        "\n",
        "\n",
        "Since all our operators are unitary, we can easily **undo** or **erase** them by applying their adjoint\n",
        "\n",
        "$$U^{\\dagger} U | \\phi \\rangle = |\\phi\\rangle.$$\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "The **adjoint differentiation method** takes advantage of the ability to erase, creating a time- and memory- efficient method for computing quantum gradients on state vector simulators. \n",
        "\n",
        "\n",
        "--------\n",
        "\n",
        "~\n",
        "[Efficient calculation of gradients in classical simulations of variational quantum algorithm](https://arxiv.org/abs/2009.02823) \n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SGNOO8LNsLMY"
      },
      "source": [
        "## Time for some code\n",
        "\n",
        "Let's explore the above equations and their implementations in a little bit more detail.\n",
        "\n",
        "\n",
        "Start by importing pennylane and pennylane's numpy"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "id": "JbNZSRrMsWs-"
      },
      "outputs": [],
      "source": [
        "import pennylane as qml\n",
        "import jax\n",
        "from jax import numpy as np\n",
        "\n",
        "jax.config.update(\"jax_platform_name\", \"cpu\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rfEO4ay9sjzx"
      },
      "source": [
        "Need a quantum circuit to simulate"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "id": "uCfQZBshsbT4"
      },
      "outputs": [],
      "source": [
        "dev = qml.device('default.qubit', wires=2)\n",
        "\n",
        "x = np.array([0.1, 0.2, 0.3])\n",
        "\n",
        "@qml.qnode(dev, diff_method=\"adjoint\")\n",
        "def circuit(a):\n",
        "    qml.RX(a[0], wires=0)\n",
        "    qml.CNOT(wires=(0,1))\n",
        "    qml.RY(a[1], wires=1)\n",
        "    qml.RZ(a[2], wires=1)\n",
        "    return qml.expval(qml.PauliX(wires=1))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LYSKH7MPsqHO"
      },
      "source": [
        "The fast c++ simulator device `\"lightning.qubit\"` also supports adjoint differentiation, but here we want to quickly prototype a minimal version to illustrate how the algorithm works. We recommend performing adjoint differentiation on `\"lightning.qubit\"` for substantial performance increases.\n",
        "\n",
        "We will use the `circuit` QNode just for comparison purposes. Throughout this demo, we will instead use a list of its operations `ops` and a single observable `M`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "id": "y2zE9-xZsnnE"
      },
      "outputs": [],
      "source": [
        "n_gates = 4\n",
        "n_params = 3\n",
        "\n",
        "ops = [\n",
        "    qml.RX(x[0], wires=0),\n",
        "    qml.CNOT(wires=(0,1)),\n",
        "    qml.RY(x[1], wires=1),\n",
        "    qml.RZ(x[2], wires=1)\n",
        "]\n",
        "M = qml.PauliX(wires=1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TMXACcejtNL_"
      },
      "source": [
        "We create our state by using the `\"default.qubit\"` methods `_create_basis_state` and `_apply_operation`.\n",
        "\n",
        "These are private methods that you shouldn’t typically use and are subject to change without a deprecation period, but we use them here to illustrate the algorithm.\n",
        "\n",
        "Internally, the device uses a 2x2x2x… array to represent the state, whereas the measurement `qml.state()` and the device attribute `dev.state` flatten this internal representation."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "vQ07KCzOsx3Z",
        "outputId": "36f8b0c7-2cea-4eba-e934-b80c05ee0cbd"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "[[9.8260182e-01-0.14850575j 9.8589033e-02+0.01490028j]\n",
            " [7.4563531e-04+0.00493356j 7.4314815e-03-0.04917107j]]\n"
          ]
        }
      ],
      "source": [
        "from pennylane.devices.qubit import create_initial_state, apply_operation\n",
        "\n",
        "state = create_initial_state((0, 1))\n",
        "\n",
        "for op in ops:\n",
        "    state = apply_operation(op, state)\n",
        "\n",
        "print(state)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "35NeH5VitXwe"
      },
      "source": [
        "We can think of the expectation $⟨M⟩$ as an inner product between bra and a ket:\n",
        "\n",
        "$$\\langle M \\rangle = \\langle b | k \\rangle = \\langle \\Psi | M | \\Psi \\rangle,$$\n",
        "\n",
        "where \n",
        "\n",
        "$$\\langle b | = \\langle \\Psi| M = \\langle 0 | U_0^{\\dagger} \\dots U_n^{\\dagger} M,$$\n",
        "\n",
        "$$| k \\rangle =  |\\Psi \\rangle = U_n U_{n-1} \\dots U_0 |0\\rangle.$$\n",
        "\n",
        "\n",
        "Using the `state` calculated above, we can create these $|b⟩$ and |k⟩$ vectors\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "id": "bar2L7CMtWSR"
      },
      "outputs": [],
      "source": [
        "bra = apply_operation(M, state)\n",
        "ket = state"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SvCbcrUuuD5p"
      },
      "source": [
        "Now we use `np.vdot` to take their inner product. `np.vdot` sums over all dimensions and takes the complex conjugate of the first input."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "KsnkvEbcuCBO",
        "outputId": "f9a858a8-66b1-446c-e56c-db84909eaf9a"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "vdot  :  (0.18884787+0j)\n",
            "QNode :  0.18884787\n"
          ]
        }
      ],
      "source": [
        "M_expval = np.vdot(bra, ket)\n",
        "print(\"vdot  : \", M_expval)\n",
        "print(\"QNode : \", circuit(x))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JTIDkinwuJ85"
      },
      "source": [
        "We got the same result via both methods! This validates our use of vdot and device methods.\n",
        "\n",
        "But the dividing line between what makes the “bra” and “ket” vector is actually fairly arbitrary. We can divide the two vectors at any point from one $⟨0|$ to the other \n",
        "$|\n",
        "0\n",
        "⟩$\n",
        ". For example, we could have used:\n",
        "\n",
        "$$\\langle b_n | = \\langle 0 | U_1^{\\dagger} \\dots  U_n^{\\dagger} M U_n,$$\n",
        "\n",
        "$$|k_n \\rangle = U_{n-1} \\dots U_1 |0\\rangle,$$\n",
        "\n",
        "and gotten the exact same results. Here, the subscript $n$\n",
        " is used to indicate that \n",
        "$U_n$\n",
        " was moved to the bra side of the expression. Let’s calculate that instead\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "QDogUYp2uH-R",
        "outputId": "f132833f-4b39-443a-a86e-de2142ca4bee"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "(0.18884787-1.8626451e-09j)\n"
          ]
        }
      ],
      "source": [
        "bra_n = create_initial_state((0, 1))\n",
        "\n",
        "for op in ops:\n",
        "    bra_n = apply_operation(op, bra_n)\n",
        "bra_n = apply_operation(M, bra_n)\n",
        "bra_n = apply_operation(qml.adjoint(ops[-1]), bra_n)\n",
        "\n",
        "ket_n = create_initial_state((0, 1))\n",
        "\n",
        "for op in ops[:-1]: # don't apply last operation\n",
        "    ket_n = apply_operation(op, ket_n)\n",
        "\n",
        "M_expval_n = np.vdot(bra_n, ket_n)\n",
        "print(M_expval_n)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "D9BUgdQ4ucWA"
      },
      "source": [
        "Same answer!\n",
        "\n",
        "We can calculate this in a more efficient way if we already have the initial `state` $\n",
        "|\n",
        "Ψ\n",
        "⟩$\n",
        ". To shift the splitting point, we don’t have to recalculate everything from scratch. We just remove the operation from the ket and add it to the bra:\n",
        "\n",
        "$$\\langle b_n | = \\langle b | U_n,$$\n",
        "\n",
        "$$|k_n\\rangle = U_n^{\\dagger} |k\\rangle .$$\n",
        "\n",
        "\n",
        "For the ket vector, you can think of $U^{\\dagger}_n$ as “eating” its corresponding unitary from the vector, erasing it from the list of operations.\n",
        "\n",
        "Of course, we actually work with the conjugate transpose of \n",
        "$⟨\n",
        "b_\n",
        "n\n",
        "|$\n",
        ",\n",
        "\n",
        "\n",
        "$$|b_n\\rangle = U_n^{\\dagger} | b \\rangle. $$\n",
        "\n",
        "Once we write it in this form, we see that the adjoint of the operation \n",
        "$U^\n",
        "†_\n",
        "n$\n",
        " operates on both \n",
        "$|\n",
        "k_\n",
        "n\n",
        "⟩$\n",
        " and \n",
        "$|\n",
        "b_\n",
        "n\n",
        "⟩$\n",
        " to move the splitting point right.\n",
        "\n",
        "Let’s call this the “version 2” method.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "RM0scC-auaXV",
        "outputId": "989e3b06-4fb7-4619-c9f1-583f768f8849"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "(0.18884787-1.869921e-09j)\n"
          ]
        }
      ],
      "source": [
        "bra_n_v2 = apply_operation(M, state)\n",
        "ket_n_v2 = state\n",
        "\n",
        "adj_op = qml.adjoint(ops[-1])\n",
        "\n",
        "bra_n_v2 = apply_operation(adj_op, bra_n_v2)\n",
        "ket_n_v2 = apply_operation(adj_op, ket_n_v2)\n",
        "\n",
        "M_expval_n_v2 = np.vdot(bra_n_v2, ket_n_v2)\n",
        "print(M_expval_n_v2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FkzP1f6dvBZb"
      },
      "source": [
        "Much simpler!\n",
        "\n",
        "We can easily iterate over all the operations to show that the same result occurs no matter where you split the operations:\n",
        "$$\\langle b_i | = \\langle b_{i+1}| U_{i},$$\n",
        "\n",
        "$$|k_{i+1} \\rangle = U_{i} |k_{i}\\rangle.$$\n",
        "\n",
        "Rewritten, we have our iteration formulas\n",
        "\n",
        "$$| b_i \\rangle = U_i^{\\dagger} |b_{i+1}\\rangle,$$\n",
        "\n",
        "$$| k_i \\rangle  = U_i^{\\dagger} |k_{i+1}\\rangle.$$\n",
        "\n",
        "For each iteration, we move an operation from the ket side to the bra side. We start near the center at \n",
        "$U_\n",
        "n$\n",
        " and reverse through the operations list until we reach \n",
        "$U_0$\n",
        ".\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "uqdlFAqku_vL",
        "outputId": "26de66e7-fb6d-4a03-e744-7fd51788b55d"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "(0.18884787-1.869921e-09j)\n",
            "(0.18884788+8.6925944e-10j)\n",
            "(0.18884788+8.6925944e-10j)\n",
            "(0.1888479+8.6925944e-10j)\n"
          ]
        }
      ],
      "source": [
        "bra_loop = apply_operation(M, state)\n",
        "ket_loop = state\n",
        "\n",
        "for op in reversed(ops):\n",
        "    adj_op = qml.adjoint(op)\n",
        "    bra_loop = apply_operation(adj_op, bra_loop)\n",
        "    ket_loop = apply_operation(adj_op, ket_loop)\n",
        "    print(np.vdot(bra_loop, ket_loop))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "m16Shqcwva2P"
      },
      "source": [
        "# Finally the Derivatives\n",
        "\n",
        "We showed how to calculate the same thing a bunch of different ways. Why is this useful? Wherever we cut, we can stick additional things in the middle. What are we sticking in the middle? The derivative of a gate.\n",
        "\n",
        "For simplicity’s sake, assume each unitary operation \n",
        "$U_\n",
        "i$\n",
        " is a function of a single parameter \n",
        "$θ_\n",
        "i$\n",
        ". For non-parametrized gates like CNOT, we say its derivative is zero. We can also generalize the algorithm to multi-parameter gates, but we leave those out for now.\n",
        "\n",
        "Remember that each parameter occurs twice in \n",
        "$⟨\n",
        "M\n",
        "⟩$\n",
        ": once in the bra and once in the ket. Therefore, we use the product rule to take the derivative with respect to both locations:\n",
        "\n",
        "\n",
        "$$\\frac{\\partial \\langle M \\rangle}{\\partial \\theta_i} =\n",
        "\\langle 0 | U_1^{\\dagger} \\dots \\frac{\\text{d} U_i^{\\dagger}}{\\text{d} \\theta_i} \\dots M \\dots U_i \\dots U_1 | 0\\rangle $$\n",
        "$$+ \\langle 0 | U_1^{\\dagger} \\dots U_i^{\\dagger} \\dots M \\dots \\frac{\\text{d} U_i}{\\text{d} \\theta_i}  \\dots U_1 |0\\rangle$$\n",
        "\n",
        "\n",
        "We can now notice that those two components are complex conjugates of each other, so we can further simplify. Note that each term is not an expectation value of a Hermitian observable, and therefore not guaranteed to be real. When we add them together, the imaginary part cancels out, and we obtain twice the value of the real part:\n",
        "\n",
        "\n",
        "$$= 2 \\cdot \\text{Re}\\left( \\langle 0 | U_1^{\\dagger} \\dots U_i^{\\dagger} \\dots M \\dots \\frac{\\text{d} U_i}{\\text{d} \\theta_i}  \\dots U_1 |0\\rangle \\right).$$\n",
        "\n",
        "\n",
        "We can take that formula and break it into its “bra” and “ket” halves for a derivative at the \n",
        "$i^{th}$ position\n",
        "\n",
        "\n",
        "$$\\frac{\\partial \\langle M \\rangle }{\\partial \\theta_i } =\n",
        "2 \\text{Re} \\left( \\langle b_i | \\frac{\\text{d} U_i }{\\text{d} \\theta_i} | k_i \\rangle \\right)$$\n",
        "\n",
        "where \n",
        "\n",
        "$$\\langle b_i | = \\langle 0 | U_1^{\\dagger} \\dots U_n^{\\dagger} M U_n \\dots U_{i+1},$$\n",
        "\n",
        "$$|k_i \\rangle = U_{i-1} \\dots U_1 |0\\rangle.$$\n",
        "\n",
        "\n",
        "Notice that \n",
        "$U_\n",
        "i$\n",
        " does not appear in either the bra or the ket in the above equations. These formulas differ from the ones we used when just calculating the expectation value. For the actual derivative calculation, we use a temporary version of the bra,\n",
        "\n",
        "\n",
        " $$| \\tilde{k}_i \\rangle = \\frac{\\text{d} U_i}{\\text{d} \\theta_i} | k_i \\rangle$$\n",
        "\n",
        " and use to get the derivatives\n",
        "\n",
        " $$\\frac{\\partial \\langle M \\rangle}{\\partial \\theta_i} = 2 \\text{Re}\\left( \\langle b_i | \\tilde{k}_i \\rangle \\right).$$\n",
        "\n",
        " Both the bra and the ket can be calculated recursively:\n",
        "\n",
        " $$| b_{i} \\rangle = U^{\\dagger}_{i+1} |b_{i+1}\\rangle,$$\n",
        "\n",
        " $$| k_{i} \\rangle = U^{\\dagger}_{i} |k_{i+1}\\rangle.$$\n",
        "\n",
        "We can iterate through the operations starting at \n",
        "$n$\n",
        " and ending at \n",
        "$1$\n",
        ".\n",
        "\n",
        "We do have to calculate initial state first, the “forward” pass:\n",
        "\n",
        "\n",
        "$$ |\\Psi\\rangle = U_{n} U_{n-1} \\dots U_0 |0\\rangle.$$\n",
        "\n",
        "\n",
        "Once we have that, we only have about the same amount of work to calculate all the derivatives, instead of quadratically more work.\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bQ-Czh0Rwij5"
      },
      "source": [
        "# Derivative of an Operator\n",
        "\n",
        "One final thing before we get back to coding: how do we get the derivative of an operator?\n",
        "\n",
        "Most parametrized gates can be represented in terms of Pauli Rotations, which can be written as\n",
        "$$U = e^{i c \\hat{G} \\theta}$$\n",
        "\n",
        "for a Pauli matrix\n",
        "$\\hat{G}$\n",
        ", a constant \n",
        "$c$\n",
        ", and the parameter \n",
        "$θ$\n",
        ". Thus we can easily calculate their derivatives:\n",
        "\n",
        "$$\\frac{\\text{d} U}{\\text{d} \\theta} = i c \\hat{G} e^{i c \\hat{G} \\theta} = i c \\hat{G} U .$$\n",
        "\n",
        "Luckily, PennyLane already has a built-in function for calculating this\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "_grtiqDmvZNW",
        "outputId": "da3e5190-144d-4746-bdea-427fda59b9b6"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "[[-0.02498958+0.j          0.        -0.49937513j]\n",
            " [ 0.        -0.49937513j -0.02498958+0.j        ]]\n"
          ]
        }
      ],
      "source": [
        "grad_op0 = qml.operation.operation_derivative(ops[0])\n",
        "print(grad_op0)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EnrCe2rjxB1K"
      },
      "source": [
        "Now for calculating the full derivative using the adjoint method!\n",
        "\n",
        "We loop over the reversed operations, just as before. But if the operation has a parameter, we calculate its derivative and append it to a list before moving on. Since the `operation_derivative` function spits back out a matrix instead of an operation, we have to use `dev._apply_unitary` instead to create \n",
        "$|k_i⟩$\n",
        "."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "3PIkIQZOxANo",
        "outputId": "80aa0aae-f493-4b7c-f5d2-74e3d6f0b4a1"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "our calculation:  [-0.018947992473840714, 0.931615948677063, -0.05841749906539917]\n",
            "comparison:  [-0.01894799  0.93161577 -0.0584175 ]\n"
          ]
        }
      ],
      "source": [
        "bra = apply_operation(M, state)\n",
        "ket = state\n",
        "\n",
        "grads = []\n",
        "\n",
        "for op in reversed(ops):\n",
        "    adj_op = qml.adjoint(op)\n",
        "    ket = apply_operation(adj_op, ket)\n",
        "\n",
        "    # Calculating the derivative\n",
        "    if op.num_params != 0:\n",
        "        dU = qml.operation.operation_derivative(op)\n",
        "        ket_temp = apply_operation(qml.QubitUnitary(dU, op.wires), ket)\n",
        "\n",
        "        dM = 2 * np.real(np.vdot(bra, ket_temp))\n",
        "        grads.append(dM)\n",
        "\n",
        "    bra = apply_operation(adj_op, bra)\n",
        "\n",
        "\n",
        "# Finally reverse the order of the gradients\n",
        "# since we calculated them in reverse\n",
        "grads = grads[::-1]\n",
        "\n",
        "print(\"our calculation: \", [float(grad) for grad in grads])\n",
        "\n",
        "grad_compare = jax.grad(circuit)(x)\n",
        "print(\"comparison: \", grad_compare)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vXP5rCHwxQTi"
      },
      "source": [
        "It matches!!!\n",
        "\n",
        "If you want to use adjoint differentiation without having to code up your own method that can support arbitrary circuits, you can use `diff_method=\"adjoint\"` in PennyLane with `\"default.qubit\"` or PennyLane’s fast C++ simulator `\"lightning.qubit\"`"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 19,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-18TflnixOX2",
        "outputId": "a69a28b7-844d-4bca-f965-41b9d21f02e4"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "[-0.01894799  0.93161577 -0.0584175 ]\n"
          ]
        }
      ],
      "source": [
        "dev_lightning = qml.device('lightning.qubit', wires=2)\n",
        "\n",
        "@qml.qnode(dev_lightning, diff_method=\"adjoint\")\n",
        "def circuit_adjoint(a):\n",
        "    qml.RX(a[0], wires=0)\n",
        "    qml.CNOT(wires=(0,1))\n",
        "    qml.RY(a[1], wires=1)\n",
        "    qml.RZ(a[2], wires=1)\n",
        "    return qml.expval(M)\n",
        "\n",
        "print(jax.grad(circuit_adjoint)(x))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eQB594CbxaHt"
      },
      "source": [
        "# Performance\n",
        "\n",
        "The algorithm gives us the correct answers, but is it worth using? Parameter-shift gradients require at least two executions per parameter, so that method gets more and more expensive with the size of the circuit, especially on simulators. Backpropagation demonstrates decent time scaling, but requires more and more memory as the circuit gets larger. Simulation of large circuits is already RAM-limited, and backpropagation constrains the size of possible circuits even more. PennyLane also achieves backpropagation derivatives from a Python simulator and interface-specific functions. The `\"lightning.qubit\"` device does not support backpropagation, so backpropagation derivatives lose the speedup from an optimized simulator.\n",
        "\n",
        "With adjoint differentiation on `\"lightning.qubit\"`, you can get the best of both worlds: fast and memory efficient.\n",
        "\n",
        "But how fast? The provided script [here](https://pennylane.ai/qml/demos/adjoint_diff_benchmarking.html) generated the following images on a mid-range laptop. The backpropagation times were produced with the Python simulator `\"default.qubit\"`, while parameter-shift and adjoint differentiation times were calculated with `\"lightning.qubit\"`. The adjoint method clearly wins out for performance\n",
        "\n",
        "\n",
        "\n",
        "![alt text](https://pennylane.ai/_images/scaling.png)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OpLtEKSCxsQz"
      },
      "source": [
        "# Conclusion\n",
        "\n",
        "So what have we learned? Adjoint differentiation is an efficient method for differentiating quantum circuits with state vector simulation. It scales nicely in time without excessive memory requirements. Now that you’ve seen how the algorithm works, you can better understand what is happening when you select adjoint differentiation from one of PennyLane’s simulators."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "authorship_tag": "ABX9TyMIQSRcOOW0Zspn3tcJi8uH",
      "include_colab_link": true,
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "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.12.2"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}