docs/source/whitepaper/6-6-advanced-automatic-differentiation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8218bf7a",
   "metadata": {},
   "source": [
    "# Advanced Automatic Differentiation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd278af7",
   "metadata": {},
   "source": [
    "## Overview\n",
    "\n",
    "In this section, we review some advanced AD tricks, especially their application to circuit simulations. With these advanced AD tricks, we can evaluate some quantum quantities more efficiently.\n",
    "\n",
    "The advanced AD is possible in TensorCircuit, as we have implemented several AD-related API in a backend agnostic way, the implementation of them closely follows the design philosophy of [jax AD implementation](https://jax.readthedocs.io/en/latest/notebooks/autodiff_cookbook.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "611eaaef",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "17ada4bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import tensorcircuit as tc\n",
    "\n",
    "K = tc.set_backend(\"tensorflow\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "da10662e",
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 6\n",
    "nlayers = 3"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0938a5f4",
   "metadata": {},
   "source": [
    "Backend agnostic AD related APIs include the following:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e47306a6",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Help on method grad in module tensorcircuit.backends.tensorflow_backend:\n",
      "\n",
      "grad(f: Callable[..., Any], argnums: Union[int, Sequence[int]] = 0, has_aux: bool = False) -> Callable[..., Any] method of tensorcircuit.backends.tensorflow_backend.TensorFlowBackend instance\n",
      "    Return the function which is the grad function of input ``f``.\n",
      "    \n",
      "    :Example:\n",
      "    \n",
      "    >>> f = lambda x,y: x**2+2*y\n",
      "    >>> g = tc.backend.grad(f)\n",
      "    >>> g(tc.num_to_tensor(1),tc.num_to_tensor(2))\n",
      "    2\n",
      "    >>> g = tc.backend.grad(f, argnums=(0,1))\n",
      "    >>> g(tc.num_to_tensor(1),tc.num_to_tensor(2))\n",
      "    [2, 2]\n",
      "    \n",
      "    :param f: the function to be differentiated\n",
      "    :type f: Callable[..., Any]\n",
      "    :param argnums: the position of args in ``f`` that are to be differentiated, defaults to be 0\n",
      "    :type argnums: Union[int, Sequence[int]], optional\n",
      "    :return: the grad function of ``f`` with the same set of arguments as ``f``\n",
      "    :rtype: Callable[..., Any]\n",
      "\n",
      "Help on method value_and_grad in module tensorcircuit.backends.tensorflow_backend:\n",
      "\n",
      "value_and_grad(f: Callable[..., Any], argnums: Union[int, Sequence[int]] = 0, has_aux: bool = False) -> Callable[..., Tuple[Any, Any]] method of tensorcircuit.backends.tensorflow_backend.TensorFlowBackend instance\n",
      "    Return the function which returns the value and grad of ``f``.\n",
      "    \n",
      "    :Example:\n",
      "    \n",
      "    >>> f = lambda x,y: x**2+2*y\n",
      "    >>> g = tc.backend.value_and_grad(f)\n",
      "    >>> g(tc.num_to_tensor(1),tc.num_to_tensor(2))\n",
      "    5, 2\n",
      "    >>> g = tc.backend.value_and_grad(f, argnums=(0,1))\n",
      "    >>> g(tc.num_to_tensor(1),tc.num_to_tensor(2))\n",
      "    5, [2, 2]\n",
      "    \n",
      "    :param f: the function to be differentiated\n",
      "    :type f: Callable[..., Any]\n",
      "    :param argnums: the position of args in ``f`` that are to be differentiated, defaults to be 0\n",
      "    :type argnums: Union[int, Sequence[int]], optional\n",
      "    :return: the value and grad function of ``f`` with the same set of arguments as ``f``\n",
      "    :rtype: Callable[..., Tuple[Any, Any]]\n",
      "\n",
      "Help on method vectorized_value_and_grad in module tensorcircuit.backends.tensorflow_backend:\n",
      "\n",
      "vectorized_value_and_grad(f: Callable[..., Any], argnums: Union[int, Sequence[int]] = 0, vectorized_argnums: Union[int, Sequence[int]] = 0, has_aux: bool = False) -> Callable[..., Tuple[Any, Any]] method of tensorcircuit.backends.tensorflow_backend.TensorFlowBackend instance\n",
      "    Return the VVAG function of ``f``. The inputs for ``f`` is (args[0], args[1], args[2], ...),\n",
      "    and the output of ``f`` is a scalar. Suppose VVAG(f) is a function with inputs in the form\n",
      "    (vargs[0], args[1], args[2], ...), where vagrs[0] has one extra dimension than args[0] in the first axis\n",
      "    and consistent with args[0] in shape for remaining dimensions, i.e. shape(vargs[0]) = [batch] + shape(args[0]).\n",
      "    (We only cover cases where ``vectorized_argnums`` defaults to 0 here for demonstration).\n",
      "    VVAG(f) returns a tuple as a value tensor with shape [batch, 1] and a gradient tuple with shape:\n",
      "    ([batch]+shape(args[argnum]) for argnum in argnums). The gradient for argnums=k is defined as\n",
      "    \n",
      "    .. math::\n",
      "    \n",
      "        g^k = \\frac{\\partial \\sum_{i\\in batch} f(vargs[0][i], args[1], ...)}{\\partial args[k]}\n",
      "    \n",
      "    Therefore, if argnums=0, the gradient is reduced to\n",
      "    \n",
      "    .. math::\n",
      "    \n",
      "        g^0_i = \\frac{\\partial f(vargs[0][i])}{\\partial vargs[0][i]}\n",
      "    \n",
      "    , which is specifically suitable for batched VQE optimization, where args[0] is the circuit parameters.\n",
      "    \n",
      "    And if argnums=1, the gradient is like\n",
      "    \n",
      "    .. math::\n",
      "        g^1_i = \\frac{\\partial \\sum_j f(vargs[0][j], args[1])}{\\partial args[1][i]}\n",
      "    \n",
      "    , which is suitable for quantum machine learning scenarios, where ``f`` is the loss function,\n",
      "    args[0] corresponds to the input data and args[1] corresponds to the weights in the QML model.\n",
      "    \n",
      "    :param f: [description]\n",
      "    :type f: Callable[..., Any]\n",
      "    :param argnums: [description], defaults to 0\n",
      "    :type argnums: Union[int, Sequence[int]], optional\n",
      "    :param vectorized_argnums: the args to be vectorized, these arguments should share the same batch shape\n",
      "        in the fist dimension\n",
      "    :type vectorized_argnums: Union[int, Sequence[int]], defaults to 0\n",
      "    :return: [description]\n",
      "    :rtype: Callable[..., Tuple[Any, Any]]\n",
      "\n",
      "Help on method vjp in module tensorcircuit.backends.tensorflow_backend:\n",
      "\n",
      "vjp(f: Callable[..., Any], inputs: Union[Any, Sequence[Any]], v: Union[Any, Sequence[Any]]) -> Tuple[Union[Any, Sequence[Any]], Union[Any, Sequence[Any]]] method of tensorcircuit.backends.tensorflow_backend.TensorFlowBackend instance\n",
      "    Function that computes the dot product between a vector v and the Jacobian\n",
      "    of the given function at the point given by the inputs. (reverse mode AD relevant)\n",
      "    Strictly speaking, this function is value_and_vjp.\n",
      "    \n",
      "    :param f: the function to carry out vjp calculation\n",
      "    :type f: Callable[..., Any]\n",
      "    :param inputs: input for ``f``\n",
      "    :type inputs: Union[Tensor, Sequence[Tensor]]\n",
      "    :param v: value vector or gradient from downstream in reverse mode AD\n",
      "        the same shape as return of function ``f``\n",
      "    :type v: Union[Tensor, Sequence[Tensor]]\n",
      "    :return: (``f(*inputs)``, vjp_tensor), where vjp_tensor is the same shape as inputs\n",
      "    :rtype: Tuple[Union[Tensor, Sequence[Tensor]], Union[Tensor, Sequence[Tensor]]]\n",
      "\n",
      "Help on method jvp in module tensorcircuit.backends.tensorflow_backend:\n",
      "\n",
      "jvp(f: Callable[..., Any], inputs: Union[Any, Sequence[Any]], v: Union[Any, Sequence[Any]]) -> Tuple[Union[Any, Sequence[Any]], Union[Any, Sequence[Any]]] method of tensorcircuit.backends.tensorflow_backend.TensorFlowBackend instance\n",
      "    Function that computes a (forward-mode) Jacobian-vector product of ``f``.\n",
      "    Strictly speaking, this function is value_and_jvp.\n",
      "    \n",
      "    :param f: The function to compute jvp\n",
      "    :type f: Callable[..., Any]\n",
      "    :param inputs: input for ``f``\n",
      "    :type inputs: Union[Tensor, Sequence[Tensor]]\n",
      "    :param v: tangents\n",
      "    :type v: Union[Tensor, Sequence[Tensor]]\n",
      "    :return: (``f(*inputs)``, jvp_tensor), where jvp_tensor is the same shape as the output of ``f``\n",
      "    :rtype: Tuple[Union[Tensor, Sequence[Tensor]], Union[Tensor, Sequence[Tensor]]]\n",
      "\n",
      "Help on method jacfwd in module tensorcircuit.backends.abstract_backend:\n",
      "\n",
      "jacfwd(f: Callable[..., Any], argnums: Union[int, Sequence[int]] = 0) -> Any method of tensorcircuit.backends.tensorflow_backend.TensorFlowBackend instance\n",
      "    Compute the Jacobian of ``f`` using the forward mode AD.\n",
      "    \n",
      "    :param f: the function whose Jacobian is required\n",
      "    :type f: Callable[..., Any]\n",
      "    :param argnums: the position of the arg as Jacobian input, defaults to 0\n",
      "    :type argnums: Union[int, Sequence[int]], optional\n",
      "    :return: outer tuple for input args, inner tuple for outputs\n",
      "    :rtype: Tensor\n",
      "\n",
      "Help on method jacrev in module tensorcircuit.backends.abstract_backend:\n",
      "\n",
      "jacrev(f: Callable[..., Any], argnums: Union[int, Sequence[int]] = 0) -> Any method of tensorcircuit.backends.tensorflow_backend.TensorFlowBackend instance\n",
      "    Compute the Jacobian of ``f`` using reverse mode AD.\n",
      "    \n",
      "    :param f: The function whose Jacobian is required\n",
      "    :type f: Callable[..., Any]\n",
      "    :param argnums: the position of the arg as Jacobian input, defaults to 0\n",
      "    :type argnums: Union[int, Sequence[int]], optional\n",
      "    :return: outer tuple for output, inner tuple for input args\n",
      "    :rtype: Tensor\n",
      "\n",
      "Help on method stop_gradient in module tensorcircuit.backends.tensorflow_backend:\n",
      "\n",
      "stop_gradient(a: Any) -> Any method of tensorcircuit.backends.tensorflow_backend.TensorFlowBackend instance\n",
      "    Stop backpropagation from ``a``.\n",
      "    \n",
      "    :param a: [description]\n",
      "    :type a: Tensor\n",
      "    :return: [description]\n",
      "    :rtype: Tensor\n",
      "\n",
      "Help on method hessian in module tensorcircuit.backends.abstract_backend:\n",
      "\n",
      "hessian(f: Callable[..., Any], argnums: Union[int, Sequence[int]] = 0) -> Any method of tensorcircuit.backends.tensorflow_backend.TensorFlowBackend instance\n",
      "\n"
     ]
    }
   ],
   "source": [
    "help(K.grad)\n",
    "help(K.value_and_grad)\n",
    "help(K.vectorized_value_and_grad)\n",
    "help(K.vjp)\n",
    "help(K.jvp)\n",
    "help(K.jacfwd)\n",
    "help(K.jacrev)\n",
    "help(K.stop_gradient)\n",
    "help(K.hessian)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e8e1aad4",
   "metadata": {},
   "source": [
    "## Forward AD\n",
    "\n",
    "Using the Jacobian vector product (``jvp``), we can compute the circuit gradient in the forward AD mode, which is more suitable when the number of output elements is much larger than the input.\n",
    "\n",
    "Suppose we are going to evaluate $\\partial \\vert \\psi(\\theta) \\rangle$, where $\\psi(\\theta) = U(\\theta)\\vert 0\\rangle$ is the output state of some parameterized quantum circuit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "4aa4895c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def ansatz(thetas):\n",
    "    c = tc.Circuit(n)\n",
    "    for j in range(nlayers):\n",
    "        for i in range(n):\n",
    "            c.rx(i, theta=thetas[j])\n",
    "        for i in range(n - 1):\n",
    "            c.cnot(i, i + 1)\n",
    "    return c\n",
    "\n",
    "\n",
    "def psi(thetas):\n",
    "    c = ansatz(thetas)\n",
    "    return c.state()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c8f728b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "state, partial_psi_partial_theta0 = K.jvp(\n",
    "    psi,\n",
    "    K.implicit_randn([nlayers]),\n",
    "    tc.array_to_tensor(np.array([1.0, 0, 0]), dtype=\"float32\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db2d1484",
   "metadata": {},
   "source": [
    "We thus obtain $\\frac{\\partial \\psi}{\\partial \\theta_0}$, since the tangent takes one in the first place and zero in other positions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "8f8ba2d7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(TensorShape([64]), TensorShape([64]))"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "state.shape, partial_psi_partial_theta0.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1fbb508d",
   "metadata": {},
   "source": [
    "## Jacobian\n",
    "\n",
    "We can compute the Jacobian row by row or col by col using vmap together with reverse mode or forward mode AD.\n",
    "\n",
    "We still use the above example, to calculate Jacobian $J_{ij}=\\frac{\\partial \\psi_i}{\\partial \\theta_j}$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "baee108c",
   "metadata": {},
   "outputs": [],
   "source": [
    "thetas = K.implicit_randn([nlayers])\n",
    "\n",
    "jac_fun = K.jit(K.jacfwd(psi))\n",
    "\n",
    "jac_value = jac_fun(thetas)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "1d666ba9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "601 µs ± 36.5 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit jac_fun(thetas)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "24c3563e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TensorShape([64, 3])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jac_value.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5a80ec5",
   "metadata": {},
   "source": [
    "We can also use reverse mode AD to obtain Jacobian."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "a23a8f39",
   "metadata": {},
   "outputs": [],
   "source": [
    "jac_fun2 = K.jit(K.jacrev(psi))\n",
    "\n",
    "jac_value2 = jac_fun2(thetas)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "e3a4ee57",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "843 µs ± 9.95 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit jac_fun2(thetas)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "a516b41e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TensorShape([64, 3])"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jac_value2.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "f7d754e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.testing.assert_allclose(np.real(jac_value), jac_value2, atol=1e-5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "85fe0dc9",
   "metadata": {},
   "source": [
    "It is worth noting that forward mode AD Jacobian is faster since the result Jacobian is a tall matrix."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0737a17",
   "metadata": {},
   "source": [
    "## Quantum Fisher Information\n",
    "\n",
    "Quantum Fisher Information is a very important quantity in quantum computation, which can be utilized in so-called quantum natural gradient descent optimization as well as variational quantum dynamics. See [reference paper](https://arxiv.org/abs/1909.02108) for more details. \n",
    "\n",
    "There are several variants of QFI like object, and the core to evaluate is all related to $\\langle \\partial \\psi \\vert \\partial \\psi\\rangle - \\langle \\partial \\psi \\vert \\psi\\rangle\\langle \\psi \\vert \\partial \\psi\\rangle$. Such quantity is easily obtained with an advanced AD framework by first getting the Jacobian for the state and then vmap the inner product over Jacobian rows. The detailed implementation can be found in the codebase ``tensorcircuit/experimental.py``. We directly call the corresponding API in this note."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "c03b20bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorcircuit.experimental import qng"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "04b88571",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:The dtype of the watched primal must be floating (e.g. tf.float32), got tf.complex64\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "TensorShape([3, 3])"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "qfi_fun = K.jit(qng(psi))\n",
    "qfi_value = qfi_fun(thetas)\n",
    "qfi_value.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "12ef1da2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "609 µs ± 14.5 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit qfi_fun(thetas) # the speed is comparable with a simple Jacobian computation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e68f4922",
   "metadata": {},
   "source": [
    "## Hessian\n",
    "\n",
    "Hessian is defined as $\\partial_{ij} \\langle \\psi(\\theta)\\vert H\\vert \\psi(\\theta)\\rangle$, where $ij$ is shorthand for $\\theta_i\\theta_j$.\n",
    "\n",
    "In the following examples, we use $H=Z_0$ for simplicity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "238c3fbc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def h(thetas):\n",
    "    c = ansatz(thetas)\n",
    "    return K.real(c.expectation_ps(z=[0]))\n",
    "\n",
    "\n",
    "hess_f = K.jit(K.hessian(h))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "62ead94a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TensorShape([3, 3])"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hess_value = hess_f(thetas)\n",
    "hess_value.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "484b1b1a",
   "metadata": {},
   "source": [
    "## $\\langle \\psi \\vert H \\vert \\partial \\psi \\rangle$\n",
    "\n",
    "This quantity is very common as the RHS of the variational quantum dynamics equation. And there is no good way to compute this quantity besides constructing a corresponding Hadamard test circuit.\n",
    "\n",
    "However, we can easily obtain this in the AD framework, as long as the ``stop_gradint`` API exists, which is the case for TensorCircuit. Namely, this quantity is obtained as $\\partial (\\langle \\psi \\vert H\\vert \\bot(\\psi)\\rangle)$, where the outside $\\partial$ is automatically implemented by AD and $\\bot$ is for ``stop_gradient`` op which stop the backpropagation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "31e21574",
   "metadata": {},
   "outputs": [],
   "source": [
    "z0 = tc.quantum.PauliStringSum2Dense([[3, 0, 0, 0, 0, 0]])\n",
    "\n",
    "\n",
    "def h(thetas):\n",
    "    w = psi(thetas)\n",
    "    wr = K.stop_gradient(w)\n",
    "    wl = K.conj(w)\n",
    "    wl = K.reshape(wl, [1, -1])\n",
    "    wr = K.reshape(wr, [-1, 1])\n",
    "    e = wl @ z0 @ wr\n",
    "    return K.real(e)[0, 0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "c3b7b267",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TensorShape([3])"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "psi_h_partial_psi = K.grad(h)(thetas)\n",
    "psi_h_partial_psi.shape"
   ]
  }
 ],
 "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.8.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}