3.3.2-Compilation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3.3.2. Augmenting differentiable compilation transforms with JIT compilation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook contains all the results for the JIT compilation examples. Note that the runtimes will vary depending on your processor speed and so will not match those in the paper exactly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "import pennylane as qml\n",
    "from pennylane import numpy as np\n",
    "from pennylane.transforms import (\n",
    "    commute_controlled,\n",
    "    single_qubit_fusion\n",
    ")\n",
    "\n",
    "import jax\n",
    "from jax import numpy as jnp\n",
    "\n",
    "# Enable float64 support\n",
    "from jax.config import config\n",
    "\n",
    "remember = config.read(\"jax_enable_x64\")\n",
    "config.update(\"jax_enable_x64\", True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we show how a compilation pipeline can be defined and applied to a circuit using the `@qml.compile` decorator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "pipeline = [\n",
    "    commute_controlled(direction='left'),\n",
    "    single_qubit_fusion\n",
    "]\n",
    "\n",
    "dev = qml.device('default.qubit', wires=3)\n",
    "\n",
    "@qml.qnode(dev)\n",
    "@qml.compile(pipeline=pipeline)\n",
    "def circuit(x, y, z):\n",
    "    qml.CNOT(wires=[0, 1])\n",
    "    qml.RX(x, wires=1)\n",
    "    qml.RY(y, wires=1)\n",
    "    qml.S(wires=1)\n",
    "    qml.CNOT(wires=[1, 2])\n",
    "    qml.Hadamard(wires=2)\n",
    "    qml.CNOT(wires=[2, 0])\n",
    "    qml.RZ(z, wires=2)\n",
    "    return qml.expval(qml.PauliZ(1))\n",
    "\n",
    "params = np.array([0.1, 0.2, 0.3], requires_grad=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0: ───────────────────────╭●──────────────────────────────────────────────╭X─┤     \n",
      "1: ──Rot(1.57,0.10,-1.57)─╰X──Rot(0.00,0.20,1.57)─╭●──────────────────────│──┤  <Z>\n",
      "2: ───────────────────────────────────────────────╰X──Rot(3.14,1.57,0.30)─╰●─┤     \n"
     ]
    }
   ],
   "source": [
    "print(qml.draw(circuit)(*params))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can compute the gradient of the three input parameters using the `qml.grad` transform."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array(-0.0978434), array(-0.19767681), array(1.33356867e-17))"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "qml.grad(circuit)(*params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# JIT compilation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we will incorporate JIT into the compilation process. Let's define the circuit structure below with a number of layers of parametrized gates. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "dev = qml.device('default.qubit', wires=5)\n",
    "\n",
    "def circuit(x, weights):\n",
    "    for wire in range(5):\n",
    "        qml.RX(x[wire], wires=wire)\n",
    "        qml.Hadamard(wires=wire)\n",
    "        \n",
    "    for wire in range(5):\n",
    "        qml.Rot(*weights[wire, :], wires=wire)\n",
    "    \n",
    "    for wire in range(5):\n",
    "        qml.CNOT(wires=[wire, (wire + 1) % 5])\n",
    "    \n",
    "    return qml.expval(\n",
    "        qml.PauliY(0) @ qml.PauliY(1) @ qml.PauliY(2) @ qml.PauliY(3) @ qml.PauliY(4)\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "original_qnode = qml.QNode(circuit, dev, interface=\"jax\", diff_method=\"parameter-shift\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These are the weights that were used to generate the example in the paper."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:jax._src.lib.xla_bridge:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)\n"
     ]
    }
   ],
   "source": [
    "x = jnp.array([0.1, 0.2, 0.3, 0.4, 0.5])\n",
    "\n",
    "weights = jnp.array([\n",
    "    [-0.28371043,  0.93681631, -1.00500712],\n",
    "    [ 1.41650132,  1.05433029,  0.91081303],\n",
    "    [-0.42656701,  0.98618842, -0.55753227],\n",
    "    [ 0.01532506, -2.07856628,  0.55483725],\n",
    "    [ 0.91423682,  0.57445956,  0.72278638]]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray(0.47543957, dtype=float64)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "original_qnode(x, weights)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0: ──RX(0.10)──H──Rot(-0.28,0.94,-1.01)─╭●──────────╭X─┤ ╭<Y@Y@Y@Y@Y>\n",
      "1: ──RX(0.20)──H──Rot(1.42,1.05,0.91)───╰X─╭●───────│──┤ ├<Y@Y@Y@Y@Y>\n",
      "2: ──RX(0.30)──H──Rot(-0.43,0.99,-0.56)────╰X─╭●────│──┤ ├<Y@Y@Y@Y@Y>\n",
      "3: ──RX(0.40)──H──Rot(0.02,-2.08,0.55)────────╰X─╭●─│──┤ ├<Y@Y@Y@Y@Y>\n",
      "4: ──RX(0.50)──H──Rot(0.91,0.57,0.72)────────────╰X─╰●─┤ ╰<Y@Y@Y@Y@Y>\n"
     ]
    }
   ],
   "source": [
    "print(qml.draw(original_qnode, expansion_strategy=\"device\")(x, weights))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "with qml.Tracker(dev) as tracker:\n",
    "    jax.grad(original_qnode, argnums=(0, 1))(x, weights)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'executions': 41, 'batches': 2, 'batch_len': 41}"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tracker.totals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "79.3 ms ± 2.44 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit \n",
    "jax.grad(original_qnode, argnums=(0, 1))(x, weights)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Now run through compiler\n",
    "compiled_qnode = qml.QNode(\n",
    "    qml.transforms.single_qubit_fusion()(circuit), \n",
    "    dev, \n",
    "    interface=\"jax\", \n",
    "    diff_method=\"parameter-shift\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "with qml.Tracker(dev) as tracker:\n",
    "    jax.grad(compiled_qnode, argnums=(0, 1))(x, weights)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'executions': 31, 'batches': 2, 'batch_len': 31}"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tracker.totals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "601 ms ± 14.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "jax.grad(compiled_qnode, argnums=(0, 1))(x, weights)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## JIT of grad"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's apply JIT to the gradient computation process of the original QNode."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "original_grad = jax.grad(original_qnode, argnums=(0, 1))\n",
    "compiled_grad = jax.grad(compiled_qnode, argnums=(0, 1))\n",
    "jitted_original_grad = jax.jit(original_grad)\n",
    "jitted_compiled_grad = jax.jit(compiled_grad)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Time the first execution manually (it will be the longest), then benchmark with JIT."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.1258864402770996\n"
     ]
    }
   ],
   "source": [
    "t0 = time.time()\n",
    "jitted_original_grad(x, weights)\n",
    "t1 = time.time()\n",
    "print(t1 - t0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "13.6 ms ± 109 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "jitted_original_grad(x, weights)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's now apply JIT to gradient computation of the compiled QNode. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "5.505212306976318\n"
     ]
    }
   ],
   "source": [
    "t0 = time.time()\n",
    "jitted_compiled_grad(x, weights)\n",
    "t1 = time.time()\n",
    "print(t1 - t0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "8.28 ms ± 95 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "jitted_compiled_grad(x, weights)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The initial run takes substantially longer, but we see great improvement afterwards."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Grad of JIT"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also reverse this process: JIT the QNode first, then compute the gradient of that."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "jitted_original_qnode = jax.jit(original_qnode)\n",
    "grad_original_jit = jax.grad(jitted_original_qnode, argnums=(0, 1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.15106821060180664\n"
     ]
    }
   ],
   "source": [
    "t0 = time.time()\n",
    "grad_original_jit(x, weights)\n",
    "t1 = time.time()\n",
    "print(t1 - t0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "16.3 ms ± 197 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "grad_original_jit(x, weights)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For the compiled QNode:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "jitted_compiled_qnode = jax.jit(compiled_qnode)\n",
    "grad_compiled_jit = jax.grad(jitted_compiled_qnode, argnums=(0, 1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "6.736414194107056\n"
     ]
    }
   ],
   "source": [
    "t0 = time.time()\n",
    "grad_compiled_jit(x, weights)\n",
    "t1 = time.time()\n",
    "print(t1 - t0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "24.4 ms ± 1.13 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "grad_compiled_jit(x, weights)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This generally leads to longer runtimes than applying JIT to the gradient computation directly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "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.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}