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

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8218bf7a",
   "metadata": {},
   "source": [
    "# 高级自动微分"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd278af7",
   "metadata": {},
   "source": [
    "## 概述\n",
    "\n",
    "在本节中，我们将回顾一些高级 AD 技巧，尤其是它们在电路模拟中的应用。借助这些高级的 AD 技巧，我们可以更高效地评估一些量子量。\n",
    "\n",
    "高级 AD 在 TensorCircuit 中是可能的，因为我们已经以后端不可知的方式实现了几个与 AD 相关的 API，它们的实现紧密遵循\n",
    "[JAX AD 实现](https://jax.readthedocs.io/en/latest/notebooks/autodiff_cookbook.html)的设计理念。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "611eaaef",
   "metadata": {},
   "source": [
    "## 设置"
   ]
  },
  {
   "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": [
    "与后端无关的 AD 相关 API 包括以下："
   ]
  },
  {
   "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": [
    "## 前向 AD\n",
    "\n",
    "使用雅可比向量积（``jvp``），我们可以计算前向 AD 模式下的电路梯度，这在输出元素的数量远大于输入的情况下更合适。\n",
    "\n",
    "假设我们要计算 $\\partial \\vert \\psi(\\theta) \\rangle$，其中 $\\psi(\\theta) = U(\\theta)\\vert 0\\rangle$ 是某个参数化量子电路的输出状态。"
   ]
  },
  {
   "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": [
    "因此我们得到 $\\frac{\\partial \\psi}{\\partial \\theta_0}$，因为切向量在第一个位置取 1，在其他位置取 0。"
   ]
  },
  {
   "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",
    "我们可以使用 vmap 和反向模式或前向模式 AD 逐行或逐列计算 Jacobian 行列式。\n",
    "\n",
    "我们仍然使用上面的例子，来计算 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": [
    "我们也可以使用反向模式 AD 来获得 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": [
    "值得注意的是，前向模式 AD Jacobian 更快，因为结果 Jacobian 是一个高矩阵。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0737a17",
   "metadata": {},
   "source": [
    "## 量子费舍尔信息\n",
    "\n",
    "量子Fisher信息是量子计算中一个非常重要的量，可用于所谓的量子自然梯度下降优化以及变分量子动力学。\n",
    "有关详细信息，请参阅 [参考论文](https://arxiv.org/abs/1909.02108)。\n",
    "\n",
    "类 QFI 的对象有多种变体，要评估的核心都与\n",
    "$\\langle \\partial \\psi \\vert \\partial \\psi\\rangle - \\langle \\partial \\psi \\vert \\psi\\rangle\\langle \\psi \\vert \\partial \\psi\\rangle$ 有关。\n",
    "使用高级 AD 框架很容易获得这样的数量，方法是首先获取状态的 Jacobian，然后在 Jacobian 行上进行 vmap 内积。\n",
    "详细的实现可以在代码库 ``tensorcircuit/experimental.py`` 中找到。我们在本笔记中直接调用对应的 API。"
   ]
  },
  {
   "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) # 速度与简单的 Jacobian 计算相当"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e68f4922",
   "metadata": {},
   "source": [
    "## Hessian\n",
    "\n",
    "Hessian 定义为 $\\partial_{ij} \\langle \\psi(\\theta)\\vert H\\vert \\psi(\\theta)\\rangle$，其中 $ij$ 是 $\\theta_i\\theta_j$ 的简写。\n",
    "\n",
    "在以下示例中，为简单起见，我们使用 $H=Z_0$。"
   ]
  },
  {
   "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",
    "这个量作为变分量子动力学方程的 RHS 非常常见。\n",
    "除了构建相应的 Hadamard 测试电路外，没有很好的方法来计算这个量。\n",
    "\n",
    "但是，只要存在 ``stop_gradint`` API，我们就可以在 AD 框架中轻松获取，TensorCircuit 就是这种情况。\n",
    "即这个量得到 $\\partial (\\langle \\psi \\vert H\\vert \\bot(\\psi)\\rangle)$，其中外部 $\\partial$ 由 AD 自动实现，\n",
    "$\\bot$ 为停止反向传播的 ``stop_gradient`` 操作。"
   ]
  },
  {
   "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
}