JAX_control_flow.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/jecampagne/JaxTutos/blob/main/JAX_control_flow.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "9qp3EqAr39Au"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "\n",
    "from jax import grad, jit, vmap\n",
    "from jax import jacfwd, jacrev\n",
    "\n",
    "jax.config.update(\"jax_enable_x64\", True)\n",
    "\n",
    "from functools import partial"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9u1tAPdu5g3l"
   },
   "source": [
    "# Theme: explore some JAX flux control\n",
    "\n",
    "Complementary to `JIT_fractals.ipyn` nb where `jax.lax.while` and `jax.lax.cond` are used."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xxD9HRzp4TAJ"
   },
   "source": [
    "## jax.lax.fori_loop\n",
    " the equivalen code is\n",
    "```python\n",
    "def fori_loop(start, stop, body_fun, init_val):\n",
    "  val = init_val\n",
    "  for i in range(start, stop):\n",
    "    val = body_fun(i, val)\n",
    "  return val\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "BV2wDOrW4ZhC"
   },
   "source": [
    "Basic use"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "PjREMsSm4ItG",
    "outputId": "13d5a2db-71cb-46cb-c5c5-c3dcf3c9c568"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "385 385\n"
     ]
    }
   ],
   "source": [
    "def body(i,val):\n",
    "  return val+i**2\n",
    "val_init=0\n",
    "N=10\n",
    "res = jax.lax.fori_loop(0,N+1,body,val_init)\n",
    "print(res, N*(N+1)*(2*N+1)//6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "a0NIzJi_4bgu",
    "outputId": "56d85201-fbfd-4b99-8581-64baa12ff92f"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10], dtype=int64)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "arr = jnp.arange(0,N+1)\n",
    "arr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "YFz1Scmb4jKw",
    "outputId": "3a16550a-5692-4258-b7e1-7521a31e8d00"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(Array(11, dtype=int64),\n",
       " Array(55, dtype=int64),\n",
       " Array(385, dtype=int64),\n",
       " Array(3025, dtype=int64))"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jnp.sum(arr**0), jnp.sum(arr**1), jnp.sum(arr**2), jnp.sum(arr**3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "id": "hzTvFOve4l-i"
   },
   "outputs": [],
   "source": [
    "#just a reminder that jax device array is immutable\n",
    "def body(dum,val):\n",
    "  a, b, iptr=val\n",
    "  b =  b.at[iptr].set(jnp.sum(a**iptr)) # a reminder that jax array is immutable\n",
    "  iptr += 1\n",
    "  return a, b, iptr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "jR4RcmgS4o8q",
    "outputId": "91cfa5ae-4520-44a4-92a2-f9a06750e31f"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([  11.,   55.,  385., 3025.], dtype=float64)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "b=jnp.zeros(shape=(4,))\n",
    "_,b,_ = jax.lax.fori_loop(0,len(b),body,(arr,b,0))\n",
    "b"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9IYsmRhP4v0U"
   },
   "source": [
    "## jax.lax.scan\n",
    "\n",
    "equivalen code\n",
    "\n",
    "```python\n",
    "def scan(f, init, xs, length=None):\n",
    "  if xs is None:\n",
    "    xs = [None] * length\n",
    "  carry = init\n",
    "  ys = []\n",
    "  for x in xs:\n",
    "    carry, y = f(carry, x)\n",
    "    ys.append(y)\n",
    "  return carry, np.stack(ys)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZmiPSd8Z40Dl"
   },
   "source": [
    "basic use to get the $\\sum_j a_j^i$ for different values of i."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "1BhZP_e64q00",
    "outputId": "ca7454b4-e9d1-4a99-c8b7-159f2fbb2ac6"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[  11   55  385 3025]\n"
     ]
    }
   ],
   "source": [
    "def body(a, i):\n",
    "  return a, jnp.sum(a**i)\n",
    "\n",
    "arr = jnp.arange(0,N+1)\n",
    "_, res = jax.lax.scan(body,arr,jnp.array([0,1,2,3]))\n",
    "print(res)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "y0uvACCI454E"
   },
   "source": [
    "## A minimizer where carry is updated w/o changing its structure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "id": "X0J7iGK342lS"
   },
   "outputs": [],
   "source": [
    "from jax.example_libraries.optimizers  import adam\n",
    "\n",
    "sum_abs = lambda x: jnp.sum(jnp.abs(x))\n",
    "\n",
    "def minimize(\n",
    "    fun, inputs, optimizer=sgd, schedule=1e-3, maxiter=4, gtol=1e-2,\n",
    "):\n",
    "    \"\"\"Optimizes inputs to minimize output of energy_fn\n",
    "    Args:\n",
    "        fun: maps x -> energy\n",
    "        inputs: initial inputs\n",
    "        optimizer: any jax.example_libraries.optimizers\n",
    "        schedule: learning rate schedule for optimizer\n",
    "        maxiter: maximum number of optimizer steps\n",
    "        gtol: cutoff for gradients to stop early\n",
    "    \"\"\"\n",
    "    init_optimizer, update_optimizer, get_inputs = optimizer(schedule)\n",
    "    energy_gradient = grad(lambda x: jnp.sum(fun(x)))\n",
    "    opt_state = init_optimizer(inputs)\n",
    "    gradient = energy_gradient(inputs)\n",
    "    sum_abs_gradient = sum_abs(gradient)\n",
    "    carry = (0, opt_state, gradient, sum_abs_gradient, inputs)\n",
    "\n",
    "    # jit not necesary done auto\n",
    "    def step_once(carry):\n",
    "        (step, opt_state, gradient, sum_abs_gradient, inputs) = carry\n",
    "        opt_state = update_optimizer(step, gradient, opt_state)\n",
    "        inputs = get_inputs(opt_state)\n",
    "        gradient = energy_gradient(inputs)\n",
    "        sum_abs_gradient = sum_abs(gradient)\n",
    "        carry = (step + 1, opt_state, gradient, sum_abs_gradient, inputs)\n",
    "        return carry\n",
    "\n",
    "    def noop(carry):\n",
    "        return carry\n",
    "\n",
    "    # jit not necesary done auto\n",
    "    def body_fn(carry, x):\n",
    "        sum_abs_gradient = carry[-2]\n",
    "        return jax.lax.cond(sum_abs_gradient < gtol, noop, step_once, carry), None\n",
    "\n",
    "    carry, _ = jax.lax.scan(body_fn, carry, None, length=maxiter)\n",
    "    return carry"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "id": "3BGnHXEt5QPX"
   },
   "outputs": [],
   "source": [
    "def fun(x):\n",
    "  return (x-1.)**2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "m7v2zpUQ5QlW",
    "outputId": "637d8d69-307a-474f-885b-8fe37dd4b8b0"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "optimal x= 1.0048182232226612\n"
     ]
    }
   ],
   "source": [
    "_, final_state, _, _, _ =  minimize(fun,0.,optimizer=adam, schedule=0.1, maxiter=50)\n",
    "print(\"optimal x=\",final_state.packed_state[0][0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "192y2bYL67yD"
   },
   "source": [
    "## Auto-diff and scan/fori_loop\n",
    "\n",
    "Let defines $f:\\mathbb{R}^n → \\mathbb{R}^3$ as\n",
    "$$\n",
    "f(x_1,\\dots,x_n) = \\begin{pmatrix}\n",
    "\\sum_{i=1}^n x_i \\\\\n",
    "\\sum_{i=1}^n x_i^2 \\\\\n",
    "\\sum_{i=1}^n x_i^3\n",
    "\\end{pmatrix}\n",
    "$$\n",
    "The jacobian is\n",
    "$$\n",
    "\\nabla f(x) =  \\begin{pmatrix}\n",
    "\\frac{\\partial f_1}{\\partial x_1}=1 & \\dots & \\frac{\\partial f_1}{\\partial x_n}\\\\\n",
    "\\frac{\\partial f_2}{\\partial x_1}=2x_1& \\ddots & \\vdots \\\\\n",
    "\\frac{\\partial f_3}{\\partial x_1}=3x_1^2 & \\dots & \\frac{\\partial f_3}{\\partial x_n} \\\\\n",
    "\\end{pmatrix}_{(x)}\n",
    "$$\n",
    "we can compute then\n",
    "$$\n",
    "\\begin{pmatrix}\n",
    "\\sum_{i=1}^n 1 \\\\\n",
    "\\sum_{i=1}^n 2 x_i \\\\\n",
    "\\sum_{i=1}^n 3 x_i^2\n",
    "\\end{pmatrix}_{(x)}\n",
    "$$\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "mF154IKW_un8"
   },
   "source": [
    "Using scan"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "id": "tpBtFBph5SRo"
   },
   "outputs": [],
   "source": [
    "@jit\n",
    "def f(x):\n",
    "\n",
    "  def body(a, i):\n",
    "    return a, jnp.sum(a**i)\n",
    "\n",
    "  _, res = jax.lax.scan(body,x,jnp.array([1,2,3]))\n",
    "\n",
    "  return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "id": "s0HHGGVy7Lie"
   },
   "outputs": [],
   "source": [
    "N=10\n",
    "x=jnp.arange(0.,N+1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "7GWn4xKP97Gu",
    "outputId": "af623951-9c83-459c-9fbe-a8f790df5a46"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([ 0.,  1.,  2.,  3.,  4.,  5.,  6.,  7.,  8.,  9., 10.], dtype=float64)"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "lmcvuMoOBlug",
    "outputId": "521323bc-699d-421c-93d6-c96eff60d714"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(Array([  55.,  385., 3025.], dtype=float64), 55, 385, 3025)"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f(x), N*(N+1)//2,  N*(N+1)*(2*N+1)//6, N*N*(N+1)*(N+1)//4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "onq--TuF7WBG",
    "outputId": "f1678e0e-222d-4514-857f-bac51a29c012"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(Array([  11.,  110., 1155.], dtype=float64), 11, 110, 1155)"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jnp.sum((jacfwd(f))(x), axis=1), N+1, N*(N+1),  N*(N+1)*(2*N+1)//2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "jA4OgsJW7djP",
    "outputId": "8e1b5480-30e1-490c-c9ee-05086ea55147"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([  11.,  110., 1155.], dtype=float64)"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jnp.sum((jacrev(f))(x), axis=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "nJnPjeCH_xWJ"
   },
   "source": [
    "Using a fori_loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "id": "2ThMfc849Ywp"
   },
   "outputs": [],
   "source": [
    "@jit\n",
    "def g(x):\n",
    "\n",
    "  def body(i,val):\n",
    "    v1, v2, v3, a = val\n",
    "    return v1+a[i], v2+a[i]**2, v3+a[i]**3, a\n",
    "\n",
    "  val_init=(0.,0.,0.,x)\n",
    "  N=x.shape[0]                   # <-------------- see later\n",
    "  res = jax.lax.fori_loop(0,N,body,val_init)\n",
    "  return jnp.array(list(res[:-1]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "4BhHhlOrBLHa",
    "outputId": "23521ca2-d4de-4eeb-c6d5-3d6283efe7c4"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([  55.,  385., 3025.], dtype=float64)"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "g(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "_6CL6DqRBRq1",
    "outputId": "724d9f77-912c-489e-db57-23b7dbb9f408"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([  11.,  110., 1155.], dtype=float64)"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jnp.sum((jacfwd(g))(x), axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Z5Dou8gADcnh",
    "outputId": "95624c54-b1c5-4bae-86d2-0a771c54b4be"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([  11.,  110., 1155.], dtype=float64)"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jnp.sum((jacrev(g))(x), axis=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "509U8rzZK2O0"
   },
   "source": [
    "In the above examples, the 'N' upper loop index `N=x.shape[0]` is known at compilation time because **array shapes are static**. So even is the array shape can change from one another, the compilation can be triggered, and we can use either jacfwd or jacrev.\n",
    "\n",
    "BUT there are cases that lower or upper indexes are not known at compilation time, so then use scan.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "295LonW9x5hw"
   },
   "source": [
    "# Crash study:  JIT partialization is a solution here\n",
    "\n",
    "As an exemple following the previous ones, consider that we trunc the loop by an upper bound \"N\" as followed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {
    "id": "t9Rz-yGwfidn"
   },
   "outputs": [],
   "source": [
    "#@partial(jit, static_argnums=(1,))  # you do it after experiencing the crash of jacrev\n",
    "@jit\n",
    "def h(x,N):\n",
    "  print(\"simple print: N:\",N, \"\\nx:\",x)   # with JIT you will see if N,x are traced\n",
    "  jax.debug.print(\"jax.debug N={}, x={}\",N,x) # to debug the true value\n",
    "  def body(i,val):\n",
    "    v1, v2, v3, a = val\n",
    "    return v1+a[i], v2+a[i]**2, v3+a[i]**3, a\n",
    "\n",
    "  val_init=(0.,0.,0.,x)\n",
    "  res = jax.lax.fori_loop(0,N,body,val_init)\n",
    "  return jnp.array(list(res[:-1]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "cO23IdrXgv6o"
   },
   "source": [
    "The compilation of \"h\" is ok ?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "_5nlTnIlgr9V",
    "outputId": "68e943da-4270-47a3-b4de-f7deb8cf7659"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "simple print: N: Traced<ShapedArray(int64[], weak_type=True)>with<DynamicJaxprTrace(level=1/0)> \n",
      "x: Traced<ShapedArray(float64[100])>with<DynamicJaxprTrace(level=1/0)>\n",
      "jax.debug N=10, x=[0.9518209  0.27508441 0.37430831 0.99631444 0.89226034 0.92994324\n",
      " 0.67738551 0.46743009 0.40886326 0.50148313 0.96618484 0.04025926\n",
      " 0.43281375 0.99263362 0.72509764 0.67913372 0.10744514 0.24842745\n",
      " 0.15251267 0.94366784 0.30291106 0.62088969 0.69586144 0.38601778\n",
      " 0.28509872 0.28087924 0.90072255 0.76528619 0.4038402  0.39701052\n",
      " 0.68259172 0.75366636 0.42712534 0.42262136 0.61110725 0.21040875\n",
      " 0.31764619 0.42032282 0.71943177 0.99899488 0.27521181 0.47553947\n",
      " 0.78721589 0.63646717 0.20117968 0.93452229 0.73657255 0.9131983\n",
      " 0.12068967 0.3215231  0.18641099 0.36958946 0.250322   0.49694021\n",
      " 0.65569871 0.4343261  0.71030355 0.00402745 0.72782986 0.57965603\n",
      " 0.09246748 0.88708062 0.41514719 0.33561038 0.67008137 0.49665383\n",
      " 0.07450867 0.91085426 0.97762343 0.41555018 0.69853759 0.22560983\n",
      " 0.72575165 0.47613287 0.82625235 0.20988516 0.26390601 0.02535034\n",
      " 0.94675827 0.214524   0.62118562 0.7985099  0.0171616  0.25542527\n",
      " 0.15455303 0.10642191 0.90699606 0.19976682 0.68087741 0.30557509\n",
      " 0.00276412 0.9355432  0.60543368 0.76485782 0.31961131 0.41300814\n",
      " 0.60144344 0.56947217 0.61354185 0.90801052]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Array([6.47489364, 4.8713031 , 4.04653465], dtype=float64)"
      ]
     },
     "execution_count": 51,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y = jax.random.uniform(jax.random.PRNGKey(10),shape=(100,))\n",
    "h(y,10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "e94aTfyiymFs"
   },
   "source": [
    "The forward derivative wrt to x is ok ?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "0bq3djvUf3jT",
    "outputId": "833c4e26-abe2-41e6-e2dd-4eac6ad017f6"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "jax.debug N=10, x=[0.9518209  0.27508441 0.37430831 0.99631444 0.89226034 0.92994324\n",
      " 0.67738551 0.46743009 0.40886326 0.50148313 0.96618484 0.04025926\n",
      " 0.43281375 0.99263362 0.72509764 0.67913372 0.10744514 0.24842745\n",
      " 0.15251267 0.94366784 0.30291106 0.62088969 0.69586144 0.38601778\n",
      " 0.28509872 0.28087924 0.90072255 0.76528619 0.4038402  0.39701052\n",
      " 0.68259172 0.75366636 0.42712534 0.42262136 0.61110725 0.21040875\n",
      " 0.31764619 0.42032282 0.71943177 0.99899488 0.27521181 0.47553947\n",
      " 0.78721589 0.63646717 0.20117968 0.93452229 0.73657255 0.9131983\n",
      " 0.12068967 0.3215231  0.18641099 0.36958946 0.250322   0.49694021\n",
      " 0.65569871 0.4343261  0.71030355 0.00402745 0.72782986 0.57965603\n",
      " 0.09246748 0.88708062 0.41514719 0.33561038 0.67008137 0.49665383\n",
      " 0.07450867 0.91085426 0.97762343 0.41555018 0.69853759 0.22560983\n",
      " 0.72575165 0.47613287 0.82625235 0.20988516 0.26390601 0.02535034\n",
      " 0.94675827 0.214524   0.62118562 0.7985099  0.0171616  0.25542527\n",
      " 0.15455303 0.10642191 0.90699606 0.19976682 0.68087741 0.30557509\n",
      " 0.00276412 0.9355432  0.60543368 0.76485782 0.31961131 0.41300814\n",
      " 0.60144344 0.56947217 0.61354185 0.90801052]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Array([10.        , 12.94978727, 14.61390931], dtype=float64)"
      ]
     },
     "execution_count": 52,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jnp.sum((jacfwd(h))(y, 10), axis=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xdOu49gPywHh"
   },
   "source": [
    "But not the reversed derivative !"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 893
    },
    "id": "_Q7RpaekNn9o",
    "outputId": "98b2bf81-0316-4e9c-ae01-d6d004e7fc20"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "jax.debug N=10, x=[0.9518209  0.27508441 0.37430831 0.99631444 0.89226034 0.92994324\n",
      " 0.67738551 0.46743009 0.40886326 0.50148313 0.96618484 0.04025926\n",
      " 0.43281375 0.99263362 0.72509764 0.67913372 0.10744514 0.24842745\n",
      " 0.15251267 0.94366784 0.30291106 0.62088969 0.69586144 0.38601778\n",
      " 0.28509872 0.28087924 0.90072255 0.76528619 0.4038402  0.39701052\n",
      " 0.68259172 0.75366636 0.42712534 0.42262136 0.61110725 0.21040875\n",
      " 0.31764619 0.42032282 0.71943177 0.99899488 0.27521181 0.47553947\n",
      " 0.78721589 0.63646717 0.20117968 0.93452229 0.73657255 0.9131983\n",
      " 0.12068967 0.3215231  0.18641099 0.36958946 0.250322   0.49694021\n",
      " 0.65569871 0.4343261  0.71030355 0.00402745 0.72782986 0.57965603\n",
      " 0.09246748 0.88708062 0.41514719 0.33561038 0.67008137 0.49665383\n",
      " 0.07450867 0.91085426 0.97762343 0.41555018 0.69853759 0.22560983\n",
      " 0.72575165 0.47613287 0.82625235 0.20988516 0.26390601 0.02535034\n",
      " 0.94675827 0.214524   0.62118562 0.7985099  0.0171616  0.25542527\n",
      " 0.15455303 0.10642191 0.90699606 0.19976682 0.68087741 0.30557509\n",
      " 0.00276412 0.9355432  0.60543368 0.76485782 0.31961131 0.41300814\n",
      " 0.60144344 0.56947217 0.61354185 0.90801052]\n"
     ]
    },
    {
     "ename": "ValueError",
     "evalue": "Reverse-mode differentiation does not work for lax.while_loop or lax.fori_loop with dynamic start/stop values. Try using lax.scan, or using fori_loop with static start/stop.",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mJaxStackTraceBeforeTransformation\u001b[0m         Traceback (most recent call last)",
      "\u001b[0;32m/usr/lib/python3.10/runpy.py\u001b[0m in \u001b[0;36m_run_module_as_main\u001b[0;34m()\u001b[0m\n\u001b[1;32m    195\u001b[0m         \u001b[0msys\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0margv\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmod_spec\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0morigin\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 196\u001b[0;31m     return _run_code(code, main_globals, None,\n\u001b[0m\u001b[1;32m    197\u001b[0m                      \"__main__\", mod_spec)\n",
      "\u001b[0;32m/usr/lib/python3.10/runpy.py\u001b[0m in \u001b[0;36m_run_code\u001b[0;34m()\u001b[0m\n\u001b[1;32m     85\u001b[0m                        __spec__ = mod_spec)\n\u001b[0;32m---> 86\u001b[0;31m     \u001b[0mexec\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcode\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mrun_globals\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     87\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0mrun_globals\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/colab_kernel_launcher.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     36\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0m__name__\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;34m'__main__'\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 37\u001b[0;31m   \u001b[0mColabKernelApp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlaunch_instance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/traitlets/config/application.py\u001b[0m in \u001b[0;36mlaunch_instance\u001b[0;34m()\u001b[0m\n\u001b[1;32m    991\u001b[0m         \u001b[0mapp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0minitialize\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0margv\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 992\u001b[0;31m         \u001b[0mapp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstart\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    993\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/ipykernel/kernelapp.py\u001b[0m in \u001b[0;36mstart\u001b[0;34m()\u001b[0m\n\u001b[1;32m    618\u001b[0m             \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 619\u001b[0;31m                 \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mio_loop\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstart\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    620\u001b[0m             \u001b[0;32mexcept\u001b[0m \u001b[0mKeyboardInterrupt\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/tornado/platform/asyncio.py\u001b[0m in \u001b[0;36mstart\u001b[0;34m()\u001b[0m\n\u001b[1;32m    194\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mstart\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 195\u001b[0;31m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0masyncio_loop\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrun_forever\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    196\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/lib/python3.10/asyncio/base_events.py\u001b[0m in \u001b[0;36mrun_forever\u001b[0;34m()\u001b[0m\n\u001b[1;32m    602\u001b[0m             \u001b[0;32mwhile\u001b[0m \u001b[0;32mTrue\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 603\u001b[0;31m                 \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_run_once\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    604\u001b[0m                 \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_stopping\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/lib/python3.10/asyncio/base_events.py\u001b[0m in \u001b[0;36m_run_once\u001b[0;34m()\u001b[0m\n\u001b[1;32m   1908\u001b[0m             \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1909\u001b[0;31m                 \u001b[0mhandle\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_run\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1910\u001b[0m         \u001b[0mhandle\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mNone\u001b[0m  \u001b[0;31m# Needed to break cycles when an exception occurs.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/lib/python3.10/asyncio/events.py\u001b[0m in \u001b[0;36m_run\u001b[0;34m()\u001b[0m\n\u001b[1;32m     79\u001b[0m         \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 80\u001b[0;31m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_context\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrun\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_callback\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_args\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     81\u001b[0m         \u001b[0;32mexcept\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mSystemExit\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mKeyboardInterrupt\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/tornado/ioloop.py\u001b[0m in \u001b[0;36m<lambda>\u001b[0;34m()\u001b[0m\n\u001b[1;32m    684\u001b[0m             future.add_done_callback(\n\u001b[0;32m--> 685\u001b[0;31m                 \u001b[0;32mlambda\u001b[0m \u001b[0mf\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_run_callback\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfunctools\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpartial\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcallback\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfuture\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    686\u001b[0m             )\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/tornado/ioloop.py\u001b[0m in \u001b[0;36m_run_callback\u001b[0;34m()\u001b[0m\n\u001b[1;32m    737\u001b[0m         \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 738\u001b[0;31m             \u001b[0mret\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcallback\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    739\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mret\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/tornado/gen.py\u001b[0m in \u001b[0;36minner\u001b[0;34m()\u001b[0m\n\u001b[1;32m    824\u001b[0m                 \u001b[0mf\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mNone\u001b[0m  \u001b[0;31m# noqa: F841\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 825\u001b[0;31m                 \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mctx_run\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrun\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    826\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/tornado/gen.py\u001b[0m in \u001b[0;36mrun\u001b[0;34m()\u001b[0m\n\u001b[1;32m    785\u001b[0m                     \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 786\u001b[0;31m                         \u001b[0myielded\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgen\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mvalue\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    787\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/ipykernel/kernelbase.py\u001b[0m in \u001b[0;36mprocess_one\u001b[0;34m()\u001b[0m\n\u001b[1;32m    360\u001b[0m                 \u001b[0;32mreturn\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 361\u001b[0;31m         \u001b[0;32myield\u001b[0m \u001b[0mgen\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmaybe_future\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdispatch\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    362\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/tornado/gen.py\u001b[0m in \u001b[0;36mwrapper\u001b[0;34m()\u001b[0m\n\u001b[1;32m    233\u001b[0m                 \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 234\u001b[0;31m                     \u001b[0myielded\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mctx_run\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnext\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    235\u001b[0m                 \u001b[0;32mexcept\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mStopIteration\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mReturn\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/ipykernel/kernelbase.py\u001b[0m in \u001b[0;36mdispatch_shell\u001b[0;34m()\u001b[0m\n\u001b[1;32m    260\u001b[0m             \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 261\u001b[0;31m                 \u001b[0;32myield\u001b[0m \u001b[0mgen\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmaybe_future\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mhandler\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mstream\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0midents\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmsg\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    262\u001b[0m             \u001b[0;32mexcept\u001b[0m \u001b[0mException\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/tornado/gen.py\u001b[0m in \u001b[0;36mwrapper\u001b[0;34m()\u001b[0m\n\u001b[1;32m    233\u001b[0m                 \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 234\u001b[0;31m                     \u001b[0myielded\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mctx_run\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnext\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    235\u001b[0m                 \u001b[0;32mexcept\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mStopIteration\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mReturn\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/ipykernel/kernelbase.py\u001b[0m in \u001b[0;36mexecute_request\u001b[0;34m()\u001b[0m\n\u001b[1;32m    538\u001b[0m         reply_content = yield gen.maybe_future(\n\u001b[0;32m--> 539\u001b[0;31m             self.do_execute(\n\u001b[0m\u001b[1;32m    540\u001b[0m                 \u001b[0mcode\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msilent\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstore_history\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/tornado/gen.py\u001b[0m in \u001b[0;36mwrapper\u001b[0;34m()\u001b[0m\n\u001b[1;32m    233\u001b[0m                 \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 234\u001b[0;31m                     \u001b[0myielded\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mctx_run\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnext\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    235\u001b[0m                 \u001b[0;32mexcept\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mStopIteration\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mReturn\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/ipykernel/ipkernel.py\u001b[0m in \u001b[0;36mdo_execute\u001b[0;34m()\u001b[0m\n\u001b[1;32m    301\u001b[0m                 \u001b[0;31m# letting shell dispatch to loop runners\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 302\u001b[0;31m                 \u001b[0mres\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mshell\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrun_cell\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcode\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstore_history\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mstore_history\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msilent\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0msilent\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    303\u001b[0m         \u001b[0;32mfinally\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/ipykernel/zmqshell.py\u001b[0m in \u001b[0;36mrun_cell\u001b[0;34m()\u001b[0m\n\u001b[1;32m    538\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_last_traceback\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 539\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0msuper\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mZMQInteractiveShell\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrun_cell\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    540\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py\u001b[0m in \u001b[0;36mrun_cell\u001b[0;34m()\u001b[0m\n\u001b[1;32m   2974\u001b[0m         \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 2975\u001b[0;31m             result = self._run_cell(\n\u001b[0m\u001b[1;32m   2976\u001b[0m                 \u001b[0mraw_cell\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstore_history\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msilent\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mshell_futures\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcell_id\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py\u001b[0m in \u001b[0;36m_run_cell\u001b[0;34m()\u001b[0m\n\u001b[1;32m   3029\u001b[0m         \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 3030\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mrunner\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcoro\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   3031\u001b[0m         \u001b[0;32mexcept\u001b[0m \u001b[0mBaseException\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/IPython/core/async_helpers.py\u001b[0m in \u001b[0;36m_pseudo_sync_runner\u001b[0;34m()\u001b[0m\n\u001b[1;32m     77\u001b[0m     \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 78\u001b[0;31m         \u001b[0mcoro\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;32mNone\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     79\u001b[0m     \u001b[0;32mexcept\u001b[0m \u001b[0mStopIteration\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0mexc\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py\u001b[0m in \u001b[0;36mrun_cell_async\u001b[0;34m()\u001b[0m\n\u001b[1;32m   3256\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 3257\u001b[0;31m                 has_raised = await self.run_ast_nodes(code_ast.body, cell_name,\n\u001b[0m\u001b[1;32m   3258\u001b[0m                        interactivity=interactivity, compiler=compiler, result=result)\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py\u001b[0m in \u001b[0;36mrun_ast_nodes\u001b[0;34m()\u001b[0m\n\u001b[1;32m   3472\u001b[0m                         \u001b[0masy\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcompare\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcode\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 3473\u001b[0;31m                     \u001b[0;32mif\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0;32mawait\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrun_code\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcode\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m,\u001b[0m  \u001b[0masync_\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0masy\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   3474\u001b[0m                         \u001b[0;32mreturn\u001b[0m \u001b[0;32mTrue\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py\u001b[0m in \u001b[0;36mrun_code\u001b[0;34m()\u001b[0m\n\u001b[1;32m   3552\u001b[0m                 \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 3553\u001b[0;31m                     \u001b[0mexec\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcode_obj\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0muser_global_ns\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0muser_ns\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   3554\u001b[0m             \u001b[0;32mfinally\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-51-f2afa20a038e>\u001b[0m in \u001b[0;36m<cell line: 2>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0my\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mjax\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrandom\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0muniform\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mjax\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrandom\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mPRNGKey\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m10\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mshape\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m100\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m \u001b[0mh\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0my\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m10\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m<ipython-input-50-1c77203c7605>\u001b[0m in \u001b[0;36mh\u001b[0;34m()\u001b[0m\n\u001b[1;32m     10\u001b[0m   \u001b[0mval_init\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m0.\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m0.\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m0.\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 11\u001b[0;31m   \u001b[0mres\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mjax\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlax\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfori_loop\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mN\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mbody\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mval_init\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     12\u001b[0m   \u001b[0;32mreturn\u001b[0m \u001b[0mjnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0marray\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mlist\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mres\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mJaxStackTraceBeforeTransformation\u001b[0m: ValueError: Reverse-mode differentiation does not work for lax.while_loop or lax.fori_loop with dynamic start/stop values. Try using lax.scan, or using fori_loop with static start/stop.\n\nThe preceding stack trace is the source of the JAX operation that, once transformed by JAX, triggered the following exception.\n\n--------------------",
      "\nThe above exception was the direct cause of the following exception:\n",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-53-1a34602b29ba>\u001b[0m in \u001b[0;36m<cell line: 1>\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mjnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msum\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mjacrev\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mh\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0my\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m10\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0maxis\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;31m# replace the jit by partial(jit) to see the diffrence\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/jax/_src/api.py\u001b[0m in \u001b[0;36mjacfun\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    925\u001b[0m       \u001b[0my\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpullback\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0maux\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_vjp\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mf_partial\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0mdyn_args\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhas_aux\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    926\u001b[0m     \u001b[0mtree_map\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpartial\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0m_check_output_dtype_jacrev\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mholomorphic\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 927\u001b[0;31m     \u001b[0mjac\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mvmap\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpullback\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0m_std_basis\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0my\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    928\u001b[0m     \u001b[0mjac\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mjac\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0margnums\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mint\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32melse\u001b[0m \u001b[0mjac\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    929\u001b[0m     \u001b[0mexample_args\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mdyn_args\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0margnums\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mint\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32melse\u001b[0m \u001b[0mdyn_args\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "    \u001b[0;31m[... skipping hidden 17 frame]\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/jax/_src/lax/control_flow/loops.py\u001b[0m in \u001b[0;36m_while_transpose_error\u001b[0;34m(*_, **kwargs)\u001b[0m\n\u001b[1;32m   1696\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1697\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0m_while_transpose_error\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0m_\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1698\u001b[0;31m   raise ValueError(\"Reverse-mode differentiation does not work for \"\n\u001b[0m\u001b[1;32m   1699\u001b[0m                    \u001b[0;34m\"lax.while_loop or lax.fori_loop with dynamic start/stop values. \"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1700\u001b[0m                    \"Try using lax.scan, or using fori_loop with static start/stop.\")\n",
      "\u001b[0;31mValueError\u001b[0m: Reverse-mode differentiation does not work for lax.while_loop or lax.fori_loop with dynamic start/stop values. Try using lax.scan, or using fori_loop with static start/stop."
     ]
    }
   ],
   "source": [
    "jnp.sum((jacrev(h))(y, 10), axis=1) # replace the jit by partial(jit) to see the diffrence"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "sVktGwNgiObE"
   },
   "source": [
    "The reason of the crash is that the compilation JAX consider N as a \"Traced\" object which is not a static variable, so the implementation of fori_loop is a while_loop which do not accept backward diff.\n",
    "\n",
    "One possibility in this example  is to tell JIT not to trace \"N\", this is done using @partial(jit, static_argnums=(1,))\n",
    "\n",
    "You can now go to the \"h\" def and replace \"@jit\" by its partialization."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Bx18q-7VuJCX"
   },
   "source": [
    "Notice that the same partialization is needed also for a scan version\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {
    "id": "IWTe7zGTn4ji"
   },
   "outputs": [],
   "source": [
    "#@partial(jit, static_argnums=(1,))\n",
    "@jit\n",
    "def fnew(x,N):\n",
    "\n",
    "  tmp=x[:N]  # requires that N not to be traced to (see after crash)\n",
    "\n",
    "  def body(a, i):\n",
    "    return a, jnp.sum(a**i)\n",
    "\n",
    "  _, res = jax.lax.scan(body,tmp,jnp.array([1,2,3]))\n",
    "\n",
    "  return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 375
    },
    "id": "XqvlOfoBq-SM",
    "outputId": "d1adf6e6-e232-426a-c736-96562c81efde"
   },
   "outputs": [
    {
     "ename": "IndexError",
     "evalue": "Array slice indices must have static start/stop/step to be used with NumPy indexing syntax. Found slice(None, Traced<ShapedArray(int64[], weak_type=True)>with<DynamicJaxprTrace(level=1/0)>, None). To index a statically sized array at a dynamic position, try lax.dynamic_slice/dynamic_update_slice (JAX does not support dynamically sized arrays within JIT compiled functions).",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mIndexError\u001b[0m                                Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-60-e78671dd2738>\u001b[0m in \u001b[0;36m<cell line: 1>\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mfnew\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0my\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m10\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "    \u001b[0;31m[... skipping hidden 11 frame]\u001b[0m\n",
      "\u001b[0;32m<ipython-input-59-83f9f48c7fc1>\u001b[0m in \u001b[0;36mfnew\u001b[0;34m(x, N)\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mfnew\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mN\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 5\u001b[0;31m   \u001b[0mtmp\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0mN\u001b[0m\u001b[0;34m]\u001b[0m  \u001b[0;31m# requires that N not to be traced to\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      6\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      7\u001b[0m   \u001b[0;32mdef\u001b[0m \u001b[0mbody\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0ma\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mi\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/jax/_src/numpy/array_methods.py\u001b[0m in \u001b[0;36mop\u001b[0;34m(self, *args)\u001b[0m\n\u001b[1;32m    734\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0m_forward_operator_to_aval\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mname\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    735\u001b[0m   \u001b[0;32mdef\u001b[0m \u001b[0mop\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 736\u001b[0;31m     \u001b[0;32mreturn\u001b[0m \u001b[0mgetattr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0maval\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34mf\"_{name}\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    737\u001b[0m   \u001b[0;32mreturn\u001b[0m \u001b[0mop\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    738\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/jax/_src/numpy/array_methods.py\u001b[0m in \u001b[0;36m_getitem\u001b[0;34m(self, item)\u001b[0m\n\u001b[1;32m    347\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    348\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0m_getitem\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mitem\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 349\u001b[0;31m   \u001b[0;32mreturn\u001b[0m \u001b[0mlax_numpy\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_rewriting_take\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mitem\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    350\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    351\u001b[0m \u001b[0;31m# Syntactic sugar for scatter operations.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/jax/_src/numpy/lax_numpy.py\u001b[0m in \u001b[0;36m_rewriting_take\u001b[0;34m(arr, idx, indices_are_sorted, unique_indices, mode, fill_value)\u001b[0m\n\u001b[1;32m   4602\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   4603\u001b[0m   \u001b[0mtreedef\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstatic_idx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdynamic_idx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_split_index_for_jit\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0midx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0marr\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mshape\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 4604\u001b[0;31m   return _gather(arr, treedef, static_idx, dynamic_idx, indices_are_sorted,\n\u001b[0m\u001b[1;32m   4605\u001b[0m                  unique_indices, mode, fill_value)\n\u001b[1;32m   4606\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/jax/_src/numpy/lax_numpy.py\u001b[0m in \u001b[0;36m_gather\u001b[0;34m(arr, treedef, static_idx, dynamic_idx, indices_are_sorted, unique_indices, mode, fill_value)\u001b[0m\n\u001b[1;32m   4611\u001b[0m             unique_indices, mode, fill_value):\n\u001b[1;32m   4612\u001b[0m   \u001b[0midx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_merge_static_and_dynamic_indices\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtreedef\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstatic_idx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdynamic_idx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 4613\u001b[0;31m   \u001b[0mindexer\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_index_to_gather\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mshape\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0marr\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0midx\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# shared with _scatter_update\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   4614\u001b[0m   \u001b[0my\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0marr\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   4615\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/jax/_src/numpy/lax_numpy.py\u001b[0m in \u001b[0;36m_index_to_gather\u001b[0;34m(x_shape, idx, normalize_indices)\u001b[0m\n\u001b[1;32m   4854\u001b[0m                \u001b[0;34m\"dynamic_update_slice (JAX does not support dynamically sized \"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   4855\u001b[0m                \"arrays within JIT compiled functions).\")\n\u001b[0;32m-> 4856\u001b[0;31m         \u001b[0;32mraise\u001b[0m \u001b[0mIndexError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmsg\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   4857\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   4858\u001b[0m       \u001b[0mstart\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstep\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mslice_size\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_preprocess_slice\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mx_shape\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mx_axis\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mIndexError\u001b[0m: Array slice indices must have static start/stop/step to be used with NumPy indexing syntax. Found slice(None, Traced<ShapedArray(int64[], weak_type=True)>with<DynamicJaxprTrace(level=1/0)>, None). To index a statically sized array at a dynamic position, try lax.dynamic_slice/dynamic_update_slice (JAX does not support dynamically sized arrays within JIT compiled functions)."
     ]
    }
   ],
   "source": [
    "fnew(y,10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 375
    },
    "id": "fTAnMdwrso6y",
    "outputId": "644ec9e8-8842-4255-e3ba-31fa4ceca34e"
   },
   "outputs": [
    {
     "ename": "IndexError",
     "evalue": "Array slice indices must have static start/stop/step to be used with NumPy indexing syntax. Found slice(None, Traced<ShapedArray(int64[], weak_type=True)>with<DynamicJaxprTrace(level=3/0)>, None). To index a statically sized array at a dynamic position, try lax.dynamic_slice/dynamic_update_slice (JAX does not support dynamically sized arrays within JIT compiled functions).",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mIndexError\u001b[0m                                Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-61-89e3410a22f1>\u001b[0m in \u001b[0;36m<cell line: 1>\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mjnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msum\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mjacrev\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfnew\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0my\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m10\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0maxis\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/jax/_src/api.py\u001b[0m in \u001b[0;36mjacfun\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    921\u001b[0m     \u001b[0mtree_map\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpartial\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0m_check_input_dtype_jacrev\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mholomorphic\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mallow_int\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdyn_args\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    922\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mhas_aux\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 923\u001b[0;31m       \u001b[0my\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpullback\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_vjp\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mf_partial\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0mdyn_args\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    924\u001b[0m     \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    925\u001b[0m       \u001b[0my\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpullback\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0maux\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_vjp\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mf_partial\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0mdyn_args\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhas_aux\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/jax/_src/api.py\u001b[0m in \u001b[0;36m_vjp\u001b[0;34m(fun, has_aux, *primals)\u001b[0m\n\u001b[1;32m   2203\u001b[0m   \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mhas_aux\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2204\u001b[0m     \u001b[0mflat_fun\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mout_tree\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mflatten_fun_nokwargs\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfun\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0min_tree\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 2205\u001b[0;31m     \u001b[0mout_primal\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mout_vjp\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mad\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvjp\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mflat_fun\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mprimals_flat\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   2206\u001b[0m     \u001b[0mout_tree\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mout_tree\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2207\u001b[0m   \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/jax/_src/interpreters/ad.py\u001b[0m in \u001b[0;36mvjp\u001b[0;34m(traceable, primals, has_aux)\u001b[0m\n\u001b[1;32m    141\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mvjp\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtraceable\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mprimals\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhas_aux\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    142\u001b[0m   \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mhas_aux\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 143\u001b[0;31m     \u001b[0mout_primals\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpvals\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mjaxpr\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mconsts\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mlinearize\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtraceable\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0mprimals\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    144\u001b[0m   \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    145\u001b[0m     \u001b[0mout_primals\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpvals\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mjaxpr\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mconsts\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0maux\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mlinearize\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtraceable\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0mprimals\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhas_aux\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/jax/_src/interpreters/ad.py\u001b[0m in \u001b[0;36mlinearize\u001b[0;34m(traceable, *primals, **kwargs)\u001b[0m\n\u001b[1;32m    130\u001b[0m   \u001b[0m_\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0min_tree\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtree_flatten\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mprimals\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mprimals\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m{\u001b[0m\u001b[0;34m}\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    131\u001b[0m   \u001b[0mjvpfun_flat\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mout_tree\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mflatten_fun\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mjvpfun\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0min_tree\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 132\u001b[0;31m   \u001b[0mjaxpr\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mout_pvals\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mconsts\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mpe\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrace_to_jaxpr_nounits\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mjvpfun_flat\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0min_pvals\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    133\u001b[0m   \u001b[0mout_primals_pvals\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mout_tangents_pvals\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtree_unflatten\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mout_tree\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mout_pvals\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    134\u001b[0m   \u001b[0;32massert\u001b[0m \u001b[0mall\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mout_primal_pval\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mis_known\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mout_primal_pval\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mout_primals_pvals\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/jax/_src/profiler.py\u001b[0m in \u001b[0;36mwrapper\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    333\u001b[0m   \u001b[0;32mdef\u001b[0m \u001b[0mwrapper\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    334\u001b[0m     \u001b[0;32mwith\u001b[0m \u001b[0mTraceAnnotation\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mname\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mdecorator_kwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 335\u001b[0;31m       \u001b[0;32mreturn\u001b[0m \u001b[0mfunc\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    336\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0mwrapper\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    337\u001b[0m   \u001b[0;32mreturn\u001b[0m \u001b[0mwrapper\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/jax/_src/interpreters/partial_eval.py\u001b[0m in \u001b[0;36mtrace_to_jaxpr_nounits\u001b[0;34m(fun, pvals, instantiate)\u001b[0m\n\u001b[1;32m    772\u001b[0m   \u001b[0;32mwith\u001b[0m \u001b[0mcore\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnew_main\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mJaxprTrace\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mname_stack\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mcurrent_name_stack\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0mmain\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    773\u001b[0m     \u001b[0mfun\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtrace_to_subjaxpr_nounits\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfun\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmain\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minstantiate\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 774\u001b[0;31m     \u001b[0mjaxpr\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mout_pvals\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mconsts\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0menv\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mfun\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcall_wrapped\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpvals\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    775\u001b[0m     \u001b[0;32massert\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0menv\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    776\u001b[0m     \u001b[0;32mdel\u001b[0m \u001b[0mmain\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfun\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0menv\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/jax/_src/linear_util.py\u001b[0m in \u001b[0;36mcall_wrapped\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    190\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    191\u001b[0m     \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 192\u001b[0;31m       \u001b[0mans\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mf\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mdict\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mparams\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    193\u001b[0m     \u001b[0;32mexcept\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    194\u001b[0m       \u001b[0;31m# Some transformations yield from inside context managers, so we have to\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "    \u001b[0;31m[... skipping hidden 11 frame]\u001b[0m\n",
      "\u001b[0;32m<ipython-input-59-83f9f48c7fc1>\u001b[0m in \u001b[0;36mfnew\u001b[0;34m(x, N)\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mfnew\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mN\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 5\u001b[0;31m   \u001b[0mtmp\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0mN\u001b[0m\u001b[0;34m]\u001b[0m  \u001b[0;31m# requires that N not to be traced to\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      6\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      7\u001b[0m   \u001b[0;32mdef\u001b[0m \u001b[0mbody\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0ma\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mi\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/jax/_src/numpy/array_methods.py\u001b[0m in \u001b[0;36mop\u001b[0;34m(self, *args)\u001b[0m\n\u001b[1;32m    734\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0m_forward_operator_to_aval\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mname\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    735\u001b[0m   \u001b[0;32mdef\u001b[0m \u001b[0mop\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 736\u001b[0;31m     \u001b[0;32mreturn\u001b[0m \u001b[0mgetattr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0maval\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34mf\"_{name}\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    737\u001b[0m   \u001b[0;32mreturn\u001b[0m \u001b[0mop\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    738\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/jax/_src/numpy/array_methods.py\u001b[0m in \u001b[0;36m_getitem\u001b[0;34m(self, item)\u001b[0m\n\u001b[1;32m    347\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    348\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0m_getitem\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mitem\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 349\u001b[0;31m   \u001b[0;32mreturn\u001b[0m \u001b[0mlax_numpy\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_rewriting_take\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mitem\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    350\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    351\u001b[0m \u001b[0;31m# Syntactic sugar for scatter operations.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/jax/_src/numpy/lax_numpy.py\u001b[0m in \u001b[0;36m_rewriting_take\u001b[0;34m(arr, idx, indices_are_sorted, unique_indices, mode, fill_value)\u001b[0m\n\u001b[1;32m   4602\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   4603\u001b[0m   \u001b[0mtreedef\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstatic_idx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdynamic_idx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_split_index_for_jit\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0midx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0marr\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mshape\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 4604\u001b[0;31m   return _gather(arr, treedef, static_idx, dynamic_idx, indices_are_sorted,\n\u001b[0m\u001b[1;32m   4605\u001b[0m                  unique_indices, mode, fill_value)\n\u001b[1;32m   4606\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/jax/_src/numpy/lax_numpy.py\u001b[0m in \u001b[0;36m_gather\u001b[0;34m(arr, treedef, static_idx, dynamic_idx, indices_are_sorted, unique_indices, mode, fill_value)\u001b[0m\n\u001b[1;32m   4611\u001b[0m             unique_indices, mode, fill_value):\n\u001b[1;32m   4612\u001b[0m   \u001b[0midx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_merge_static_and_dynamic_indices\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtreedef\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstatic_idx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdynamic_idx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 4613\u001b[0;31m   \u001b[0mindexer\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_index_to_gather\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mshape\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0marr\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0midx\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# shared with _scatter_update\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   4614\u001b[0m   \u001b[0my\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0marr\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   4615\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/jax/_src/numpy/lax_numpy.py\u001b[0m in \u001b[0;36m_index_to_gather\u001b[0;34m(x_shape, idx, normalize_indices)\u001b[0m\n\u001b[1;32m   4854\u001b[0m                \u001b[0;34m\"dynamic_update_slice (JAX does not support dynamically sized \"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   4855\u001b[0m                \"arrays within JIT compiled functions).\")\n\u001b[0;32m-> 4856\u001b[0;31m         \u001b[0;32mraise\u001b[0m \u001b[0mIndexError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmsg\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   4857\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   4858\u001b[0m       \u001b[0mstart\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstep\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mslice_size\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_preprocess_slice\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mx_shape\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mx_axis\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mIndexError\u001b[0m: Array slice indices must have static start/stop/step to be used with NumPy indexing syntax. Found slice(None, Traced<ShapedArray(int64[], weak_type=True)>with<DynamicJaxprTrace(level=3/0)>, None). To index a statically sized array at a dynamic position, try lax.dynamic_slice/dynamic_update_slice (JAX does not support dynamically sized arrays within JIT compiled functions)."
     ]
    }
   ],
   "source": [
    "jnp.sum((jacrev(fnew))(y, 10), axis=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9XVS8ubX1EKh"
   },
   "source": [
    "# fori_loop vs scan\n",
    "\n",
    "Let use compute for 2 arrays\n",
    "$$\n",
    "S(x,y) = \\frac{1}{N} \\sum_{i=1,N} x_i y_i\n",
    "$$\n",
    "\n",
    "Of course if we do not consider auto-diff, S can be computed as\n",
    "```python\n",
    "x@y / x.shape[0]\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {
    "id": "com96dlu1-ai"
   },
   "outputs": [],
   "source": [
    "_, key1, key2 = jax.random.split(jax.random.PRNGKey(10),3)\n",
    "N = 100_000\n",
    "x = jax.random.uniform(key1, shape=(N,), dtype=jnp.float32)\n",
    "y = jax.random.uniform(key2, shape=(N,), dtype=jnp.float32)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "mtqYanc028vX",
    "outputId": "856e43b8-988a-4755-ce1d-57b2f9ad5a66"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(0.25065652, dtype=float32)"
      ]
     },
     "execution_count": 63,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x@y / x.shape[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {
    "id": "p106ibCc1Day"
   },
   "outputs": [],
   "source": [
    "@jit\n",
    "def loop_fun(x, y):\n",
    "\n",
    "    n = x.shape[0]   #<------  shape / dim / regular numpy ok to get static variables\n",
    "\n",
    "    def body(i, curr):\n",
    "        return curr + x[i] * y[i]\n",
    "\n",
    "    return jax.lax.fori_loop(0, n, body, 0.) / n\n",
    "\n",
    "@jit\n",
    "def scan_fun(x, y):\n",
    "    n = x.shape[0]  #<------------ idem\n",
    "\n",
    "    def body(carry, x):\n",
    "        curr, i = carry\n",
    "        return (curr + x * y[i], i+1), None\n",
    "\n",
    "    (res, _), _ = jax.lax.scan(body, (0., 0), x)\n",
    "    return res / n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Uj8ZHft83Af-",
    "outputId": "cb68d3d6-3f7b-43b1-f01b-d08fe1917739"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(0.2506565, dtype=float32)"
      ]
     },
     "execution_count": 67,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "loop_fun(x,y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Tq5E8_pa3ijW",
    "outputId": "c6058c04-d5f2-4bff-8d6e-5a25ab4af9b0"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(0.2506565, dtype=float32)"
      ]
     },
     "execution_count": 68,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "scan_fun(x, y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "ad56tPYwBAU9",
    "outputId": "32d6bc96-346c-4b16-ae5a-d99ea38551f9"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([8.4565936e-06, 8.8608285e-06, 2.0913146e-06, ..., 4.6913266e-07,\n",
       "       4.2381666e-06, 1.8729579e-06], dtype=float32)"
      ]
     },
     "execution_count": 70,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jax.jacrev(scan_fun)(x,y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "pH4CWuOnbia4",
    "outputId": "e9f20af5-5f34-41e7-a798-fb41f28f78bf"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(True, dtype=bool)"
      ]
     },
     "execution_count": 71,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.alltrue(jax.jacrev(scan_fun)(x,y) ==  jax.jacrev(loop_fun)(x,y))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3pcru_uiOI5a"
   },
   "source": [
    "# While: avoid as possible... crash here due to JIT or Jacobian Backward mode\n",
    "\n",
    "Think diffrently !\n",
    "\n",
    "Here an example: let\n",
    "$$\n",
    "X = (x_0,x_2,\\dots,x_{n-1})\\in [0,1]^n, \\quad \\sum_{i=0}^{n-1} x_i =  1\n",
    "$$\n",
    "Goal : set the highest value of $(x_i)_i$ to $0$ while $\\sum_{i=0}^{n-1} x_i> 0.5$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "nebYw6N40a7w"
   },
   "source": [
    "## simple while loop (JIT crash here)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "metadata": {
    "id": "kqzYMh1UV9tO"
   },
   "outputs": [],
   "source": [
    "def func(x):\n",
    "  x  = x/jnp.sum(x)    # normalize x to get unit-vector\n",
    "\n",
    "  while jnp.sum(x) > 0.5:\n",
    "    xold = x\n",
    "    idx = jax.lax.argmax(x,0,index_dtype=jnp.int16)\n",
    "    x = x.at[idx].set(0.)\n",
    "  return xold\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "JsDcPepNbD_U",
    "outputId": "186fc22c-1058-45f3-c178-c7fdb1c664ce"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "the input:  [0.63337177 0.50169736 0.65120645 0.12358142 0.96813596 0.83877037\n",
      " 0.91412169 0.60858754 0.88261905 0.71854831]\n",
      "the solution: [0.09258955 0.07334071 0.09519672 0.01806577 0.         0.12261578\n",
      " 0.         0.08896646 0.         0.1050411 ]  and constraint:  0.5958160737034848\n"
     ]
    }
   ],
   "source": [
    "x = jax.random.uniform(jax.random.PRNGKey(10), shape=(10,))\n",
    "print(\"the input: \",x)\n",
    "y = func(x)\n",
    "print(\"the solution:\", y, \" and constraint: \", jnp.sum(y))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "metadata": {
    "id": "ncDS0pA7e-W-"
   },
   "outputs": [],
   "source": [
    "jacfwd_f= jax.jacfwd(func)(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "0YgaYe1Af1ck",
    "outputId": "ed6d28db-660d-47a5-b343-04dada1a4531"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 96,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.allclose(jax.jacrev(func)(x),jacfwd_f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 97,
   "metadata": {
    "id": "j3gTvH11QxIY"
   },
   "outputs": [],
   "source": [
    "@jit\n",
    "def jfunc(x):\n",
    "  x  = x/jnp.sum(x)\n",
    "\n",
    "  while jnp.sum(x)>0.5:    #<------ this will crash\n",
    "    xold = x\n",
    "    idx = jax.lax.argmax(x,0,index_dtype=jnp.int16)\n",
    "    x = x.at[idx].set(0.)\n",
    "\n",
    "\n",
    "  return xold"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 106,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 468
    },
    "id": "D2gYaOJngl3b",
    "outputId": "4e922bdb-b3ac-4d76-c60c-7dbedd04a85d"
   },
   "outputs": [
    {
     "ename": "TracerBoolConversionError",
     "evalue": "Attempted boolean conversion of traced array with shape bool[]..\nThe error occurred while tracing the function jfunc at <ipython-input-97-2b837e284f75>:1 for jit. This concrete value was not available in Python because it depends on the value of the argument x.\nSee https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerBoolConversionError",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTracerBoolConversionError\u001b[0m                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-106-71388ec941eb>\u001b[0m in \u001b[0;36m<cell line: 2>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0mx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mjax\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrandom\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0muniform\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mjax\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrandom\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mPRNGKey\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m10\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mshape\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m10\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m \u001b[0my\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mjfunc\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      3\u001b[0m \u001b[0;31m#y, jnp.sum(y)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "    \u001b[0;31m[... skipping hidden 11 frame]\u001b[0m\n",
      "\u001b[0;32m<ipython-input-97-2b837e284f75>\u001b[0m in \u001b[0;36mjfunc\u001b[0;34m(x)\u001b[0m\n\u001b[1;32m      3\u001b[0m   \u001b[0mx\u001b[0m  \u001b[0;34m=\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m/\u001b[0m\u001b[0mjnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msum\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 5\u001b[0;31m   \u001b[0;32mwhile\u001b[0m \u001b[0mjnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msum\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m>\u001b[0m\u001b[0;36m0.5\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      6\u001b[0m     \u001b[0mxold\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      7\u001b[0m     \u001b[0midx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mjax\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlax\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0margmax\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mindex_dtype\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mjnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mint16\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "    \u001b[0;31m[... skipping hidden 1 frame]\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/jax/_src/core.py\u001b[0m in \u001b[0;36merror\u001b[0;34m(self, arg)\u001b[0m\n\u001b[1;32m   1490\u001b[0m   \u001b[0;32mif\u001b[0m \u001b[0mfun\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0mbool\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1491\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0merror\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0marg\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1492\u001b[0;31m       \u001b[0;32mraise\u001b[0m \u001b[0mTracerBoolConversionError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0marg\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1493\u001b[0m   \u001b[0;32melif\u001b[0m \u001b[0mfun\u001b[0m \u001b[0;32min\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mhex\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moct\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moperator\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mindex\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1494\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0merror\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0marg\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mTracerBoolConversionError\u001b[0m: Attempted boolean conversion of traced array with shape bool[]..\nThe error occurred while tracing the function jfunc at <ipython-input-97-2b837e284f75>:1 for jit. This concrete value was not available in Python because it depends on the value of the argument x.\nSee https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerBoolConversionError"
     ]
    }
   ],
   "source": [
    "x = jax.random.uniform(jax.random.PRNGKey(10), shape=(10,))\n",
    "y = jfunc(x)\n",
    "#y, jnp.sum(y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "SqKSK-R80jlQ"
   },
   "source": [
    "## with jax.lax.while_loop : jacrev crash here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "UdtkGJrtyHIM"
   },
   "source": [
    "```python\n",
    "def while_loop(cond_fun, body_fun, init_val):\n",
    "  val = init_val\n",
    "  while cond_fun(val):\n",
    "    val = body_fun(val)\n",
    "  return val\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 104,
   "metadata": {
    "id": "1tFeIkAqx320"
   },
   "outputs": [],
   "source": [
    "@jit\n",
    "def new_func(x):\n",
    "  x  = x/jnp.sum(x)\n",
    "\n",
    "  def cond(val):\n",
    "    sum, _, _ = val\n",
    "    # if sum>0.5: return True crash due to ConcretizationTypeError even w/o JIT\n",
    "    return jnp.where(sum>0.5, True, False)\n",
    "\n",
    "  def body(val):\n",
    "    sum, x, xold = val\n",
    "    xold = x\n",
    "    idx = jax.lax.argmax(x,0,index_dtype=jnp.int16)\n",
    "    x = x.at[idx].set(0.)\n",
    "    return jnp.sum(x), x, xold\n",
    "\n",
    "  val= jax.lax.while_loop(cond, body,(1.,x,x))\n",
    "\n",
    "  return val[2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 107,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "psCgeeh7zBxI",
    "outputId": "97ce66d0-aee7-4c12-f71d-e51f6a553627"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(Array([0.09258955, 0.07334071, 0.09519672, 0.01806577, 0.        ,\n",
       "        0.12261578, 0.        , 0.08896646, 0.        , 0.1050411 ],      dtype=float64),\n",
       " Array(0.59581607, dtype=float64))"
      ]
     },
     "execution_count": 107,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = jax.random.uniform(jax.random.PRNGKey(10), shape=(10,))\n",
    "y = new_func(x)\n",
    "y, jnp.sum(y)  # gives the same result as before"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 110,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Ykb04D8IztvP",
    "outputId": "2e7aaf70-ebd7-4d98-e2b7-913579a48495"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 110,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.allclose(jax.jacfwd(new_func)(x),jacfwd_f) # True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 111,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 579
    },
    "id": "tHBb6mVJzwz8",
    "outputId": "7f9179c3-4f4a-4ad6-be9b-20b95216aa73"
   },
   "outputs": [
    {
     "ename": "ValueError",
     "evalue": "Reverse-mode differentiation does not work for lax.while_loop or lax.fori_loop with dynamic start/stop values. Try using lax.scan, or using fori_loop with static start/stop.",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mJaxStackTraceBeforeTransformation\u001b[0m         Traceback (most recent call last)",
      "\u001b[0;32m/usr/lib/python3.10/runpy.py\u001b[0m in \u001b[0;36m_run_module_as_main\u001b[0;34m()\u001b[0m\n\u001b[1;32m    195\u001b[0m         \u001b[0msys\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0margv\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmod_spec\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0morigin\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 196\u001b[0;31m     return _run_code(code, main_globals, None,\n\u001b[0m\u001b[1;32m    197\u001b[0m                      \"__main__\", mod_spec)\n",
      "\u001b[0;32m/usr/lib/python3.10/runpy.py\u001b[0m in \u001b[0;36m_run_code\u001b[0;34m()\u001b[0m\n\u001b[1;32m     85\u001b[0m                        __spec__ = mod_spec)\n\u001b[0;32m---> 86\u001b[0;31m     \u001b[0mexec\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcode\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mrun_globals\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     87\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0mrun_globals\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/colab_kernel_launcher.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     36\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0m__name__\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;34m'__main__'\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 37\u001b[0;31m   \u001b[0mColabKernelApp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlaunch_instance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/traitlets/config/application.py\u001b[0m in \u001b[0;36mlaunch_instance\u001b[0;34m()\u001b[0m\n\u001b[1;32m    991\u001b[0m         \u001b[0mapp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0minitialize\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0margv\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 992\u001b[0;31m         \u001b[0mapp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstart\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    993\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/ipykernel/kernelapp.py\u001b[0m in \u001b[0;36mstart\u001b[0;34m()\u001b[0m\n\u001b[1;32m    618\u001b[0m             \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 619\u001b[0;31m                 \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mio_loop\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstart\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    620\u001b[0m             \u001b[0;32mexcept\u001b[0m \u001b[0mKeyboardInterrupt\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/tornado/platform/asyncio.py\u001b[0m in \u001b[0;36mstart\u001b[0;34m()\u001b[0m\n\u001b[1;32m    194\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mstart\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 195\u001b[0;31m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0masyncio_loop\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrun_forever\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    196\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/lib/python3.10/asyncio/base_events.py\u001b[0m in \u001b[0;36mrun_forever\u001b[0;34m()\u001b[0m\n\u001b[1;32m    602\u001b[0m             \u001b[0;32mwhile\u001b[0m \u001b[0;32mTrue\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 603\u001b[0;31m                 \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_run_once\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    604\u001b[0m                 \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_stopping\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/lib/python3.10/asyncio/base_events.py\u001b[0m in \u001b[0;36m_run_once\u001b[0;34m()\u001b[0m\n\u001b[1;32m   1908\u001b[0m             \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1909\u001b[0;31m                 \u001b[0mhandle\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_run\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1910\u001b[0m         \u001b[0mhandle\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mNone\u001b[0m  \u001b[0;31m# Needed to break cycles when an exception occurs.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/lib/python3.10/asyncio/events.py\u001b[0m in \u001b[0;36m_run\u001b[0;34m()\u001b[0m\n\u001b[1;32m     79\u001b[0m         \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 80\u001b[0;31m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_context\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrun\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_callback\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_args\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     81\u001b[0m         \u001b[0;32mexcept\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mSystemExit\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mKeyboardInterrupt\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/tornado/ioloop.py\u001b[0m in \u001b[0;36m<lambda>\u001b[0;34m()\u001b[0m\n\u001b[1;32m    684\u001b[0m             future.add_done_callback(\n\u001b[0;32m--> 685\u001b[0;31m                 \u001b[0;32mlambda\u001b[0m \u001b[0mf\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_run_callback\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfunctools\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpartial\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcallback\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfuture\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    686\u001b[0m             )\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/tornado/ioloop.py\u001b[0m in \u001b[0;36m_run_callback\u001b[0;34m()\u001b[0m\n\u001b[1;32m    737\u001b[0m         \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 738\u001b[0;31m             \u001b[0mret\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcallback\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    739\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mret\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/tornado/gen.py\u001b[0m in \u001b[0;36minner\u001b[0;34m()\u001b[0m\n\u001b[1;32m    824\u001b[0m                 \u001b[0mf\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mNone\u001b[0m  \u001b[0;31m# noqa: F841\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 825\u001b[0;31m                 \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mctx_run\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrun\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    826\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/tornado/gen.py\u001b[0m in \u001b[0;36mrun\u001b[0;34m()\u001b[0m\n\u001b[1;32m    785\u001b[0m                     \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 786\u001b[0;31m                         \u001b[0myielded\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgen\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mvalue\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    787\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/ipykernel/kernelbase.py\u001b[0m in \u001b[0;36mprocess_one\u001b[0;34m()\u001b[0m\n\u001b[1;32m    360\u001b[0m                 \u001b[0;32mreturn\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 361\u001b[0;31m         \u001b[0;32myield\u001b[0m \u001b[0mgen\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmaybe_future\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdispatch\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    362\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/tornado/gen.py\u001b[0m in \u001b[0;36mwrapper\u001b[0;34m()\u001b[0m\n\u001b[1;32m    233\u001b[0m                 \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 234\u001b[0;31m                     \u001b[0myielded\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mctx_run\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnext\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    235\u001b[0m                 \u001b[0;32mexcept\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mStopIteration\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mReturn\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/ipykernel/kernelbase.py\u001b[0m in \u001b[0;36mdispatch_shell\u001b[0;34m()\u001b[0m\n\u001b[1;32m    260\u001b[0m             \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 261\u001b[0;31m                 \u001b[0;32myield\u001b[0m \u001b[0mgen\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmaybe_future\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mhandler\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mstream\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0midents\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmsg\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    262\u001b[0m             \u001b[0;32mexcept\u001b[0m \u001b[0mException\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/tornado/gen.py\u001b[0m in \u001b[0;36mwrapper\u001b[0;34m()\u001b[0m\n\u001b[1;32m    233\u001b[0m                 \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 234\u001b[0;31m                     \u001b[0myielded\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mctx_run\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnext\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    235\u001b[0m                 \u001b[0;32mexcept\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mStopIteration\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mReturn\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/ipykernel/kernelbase.py\u001b[0m in \u001b[0;36mexecute_request\u001b[0;34m()\u001b[0m\n\u001b[1;32m    538\u001b[0m         reply_content = yield gen.maybe_future(\n\u001b[0;32m--> 539\u001b[0;31m             self.do_execute(\n\u001b[0m\u001b[1;32m    540\u001b[0m                 \u001b[0mcode\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msilent\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstore_history\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/tornado/gen.py\u001b[0m in \u001b[0;36mwrapper\u001b[0;34m()\u001b[0m\n\u001b[1;32m    233\u001b[0m                 \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 234\u001b[0;31m                     \u001b[0myielded\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mctx_run\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnext\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    235\u001b[0m                 \u001b[0;32mexcept\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mStopIteration\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mReturn\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/ipykernel/ipkernel.py\u001b[0m in \u001b[0;36mdo_execute\u001b[0;34m()\u001b[0m\n\u001b[1;32m    301\u001b[0m                 \u001b[0;31m# letting shell dispatch to loop runners\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 302\u001b[0;31m                 \u001b[0mres\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mshell\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrun_cell\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcode\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstore_history\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mstore_history\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msilent\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0msilent\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    303\u001b[0m         \u001b[0;32mfinally\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/ipykernel/zmqshell.py\u001b[0m in \u001b[0;36mrun_cell\u001b[0;34m()\u001b[0m\n\u001b[1;32m    538\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_last_traceback\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 539\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0msuper\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mZMQInteractiveShell\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrun_cell\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    540\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py\u001b[0m in \u001b[0;36mrun_cell\u001b[0;34m()\u001b[0m\n\u001b[1;32m   2974\u001b[0m         \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 2975\u001b[0;31m             result = self._run_cell(\n\u001b[0m\u001b[1;32m   2976\u001b[0m                 \u001b[0mraw_cell\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstore_history\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msilent\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mshell_futures\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcell_id\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py\u001b[0m in \u001b[0;36m_run_cell\u001b[0;34m()\u001b[0m\n\u001b[1;32m   3029\u001b[0m         \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 3030\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mrunner\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcoro\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   3031\u001b[0m         \u001b[0;32mexcept\u001b[0m \u001b[0mBaseException\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/IPython/core/async_helpers.py\u001b[0m in \u001b[0;36m_pseudo_sync_runner\u001b[0;34m()\u001b[0m\n\u001b[1;32m     77\u001b[0m     \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 78\u001b[0;31m         \u001b[0mcoro\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;32mNone\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     79\u001b[0m     \u001b[0;32mexcept\u001b[0m \u001b[0mStopIteration\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0mexc\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py\u001b[0m in \u001b[0;36mrun_cell_async\u001b[0;34m()\u001b[0m\n\u001b[1;32m   3256\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 3257\u001b[0;31m                 has_raised = await self.run_ast_nodes(code_ast.body, cell_name,\n\u001b[0m\u001b[1;32m   3258\u001b[0m                        interactivity=interactivity, compiler=compiler, result=result)\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py\u001b[0m in \u001b[0;36mrun_ast_nodes\u001b[0;34m()\u001b[0m\n\u001b[1;32m   3472\u001b[0m                         \u001b[0masy\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcompare\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcode\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 3473\u001b[0;31m                     \u001b[0;32mif\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0;32mawait\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrun_code\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcode\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m,\u001b[0m  \u001b[0masync_\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0masy\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   3474\u001b[0m                         \u001b[0;32mreturn\u001b[0m \u001b[0;32mTrue\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py\u001b[0m in \u001b[0;36mrun_code\u001b[0;34m()\u001b[0m\n\u001b[1;32m   3552\u001b[0m                 \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 3553\u001b[0;31m                     \u001b[0mexec\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcode_obj\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0muser_global_ns\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0muser_ns\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   3554\u001b[0m             \u001b[0;32mfinally\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-105-cfba1bdb1f0c>\u001b[0m in \u001b[0;36m<cell line: 2>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0mx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mjax\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrandom\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0muniform\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mjax\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrandom\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mPRNGKey\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m10\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mshape\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m10\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m \u001b[0my\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mnew_func\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      3\u001b[0m \u001b[0my\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mjnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msum\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0my\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-104-f9a6aab68265>\u001b[0m in \u001b[0;36mnew_func\u001b[0;34m()\u001b[0m\n\u001b[1;32m     16\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 17\u001b[0;31m   \u001b[0mval\u001b[0m\u001b[0;34m=\u001b[0m \u001b[0mjax\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlax\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mwhile_loop\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcond\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbody\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m1.\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     18\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mJaxStackTraceBeforeTransformation\u001b[0m: ValueError: Reverse-mode differentiation does not work for lax.while_loop or lax.fori_loop with dynamic start/stop values. Try using lax.scan, or using fori_loop with static start/stop.\n\nThe preceding stack trace is the source of the JAX operation that, once transformed by JAX, triggered the following exception.\n\n--------------------",
      "\nThe above exception was the direct cause of the following exception:\n",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-111-99cc50c11509>\u001b[0m in \u001b[0;36m<cell line: 1>\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mallclose\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mjax\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mjacrev\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnew_func\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mjacfwd_f\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;31m# Crash\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/jax/_src/api.py\u001b[0m in \u001b[0;36mjacfun\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    925\u001b[0m       \u001b[0my\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpullback\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0maux\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_vjp\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mf_partial\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0mdyn_args\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhas_aux\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    926\u001b[0m     \u001b[0mtree_map\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpartial\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0m_check_output_dtype_jacrev\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mholomorphic\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 927\u001b[0;31m     \u001b[0mjac\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mvmap\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpullback\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0m_std_basis\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0my\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    928\u001b[0m     \u001b[0mjac\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mjac\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0margnums\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mint\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32melse\u001b[0m \u001b[0mjac\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    929\u001b[0m     \u001b[0mexample_args\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mdyn_args\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0margnums\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mint\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32melse\u001b[0m \u001b[0mdyn_args\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "    \u001b[0;31m[... skipping hidden 17 frame]\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/jax/_src/lax/control_flow/loops.py\u001b[0m in \u001b[0;36m_while_transpose_error\u001b[0;34m(*_, **kwargs)\u001b[0m\n\u001b[1;32m   1696\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1697\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0m_while_transpose_error\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0m_\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1698\u001b[0;31m   raise ValueError(\"Reverse-mode differentiation does not work for \"\n\u001b[0m\u001b[1;32m   1699\u001b[0m                    \u001b[0;34m\"lax.while_loop or lax.fori_loop with dynamic start/stop values. \"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1700\u001b[0m                    \"Try using lax.scan, or using fori_loop with static start/stop.\")\n",
      "\u001b[0;31mValueError\u001b[0m: Reverse-mode differentiation does not work for lax.while_loop or lax.fori_loop with dynamic start/stop values. Try using lax.scan, or using fori_loop with static start/stop."
     ]
    }
   ],
   "source": [
    "np.allclose(jax.jacrev(new_func)(x),jacfwd_f) # Crash"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "y5XAjQ6az9VM"
   },
   "source": [
    "The solution: use jax.lax.cond & jax.lax.scan"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 112,
   "metadata": {
    "id": "KDtELc4PifOR"
   },
   "outputs": [],
   "source": [
    "def jfunc_bis(x):\n",
    "  n = x.size\n",
    "  x  = x/jnp.sum(x)\n",
    "\n",
    "  def noop(carry):\n",
    "    return carry\n",
    "\n",
    "  @jit\n",
    "  def step_once(carry):\n",
    "    (step, sum, x, xold) = carry\n",
    "    xold = x\n",
    "    idx = jax.lax.argmax(x,0,index_dtype=jnp.int16)\n",
    "    x = x.at[idx].set(0.)\n",
    "    carry = (step+1, jnp.sum(x), x, xold)\n",
    "    return carry\n",
    "\n",
    "  @jit\n",
    "  def body_fn(carry, dummy):\n",
    "      (step, sum, x, xold) = carry\n",
    "      return jax.lax.cond(sum < 0.5, noop, step_once, carry), None\n",
    "\n",
    "\n",
    "  carry = (0, 1., x, x)\n",
    "  carry, _ = jax.lax.scan(body_fn, carry, None, length=n)\n",
    "  return carry[3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 114,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "5SZ5Xk9tlRfc",
    "outputId": "e5deeb12-e068-4cb1-f9d2-3f1e566ea6bc"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.63337177 0.50169736 0.65120645 0.12358142 0.96813596 0.83877037\n",
      " 0.91412169 0.60858754 0.88261905 0.71854831]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(Array([0.09258955, 0.07334071, 0.09519672, 0.01806577, 0.        ,\n",
       "        0.12261578, 0.        , 0.08896646, 0.        , 0.1050411 ],      dtype=float64),\n",
       " Array(0.59581607, dtype=float64))"
      ]
     },
     "execution_count": 114,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = jax.random.uniform(jax.random.PRNGKey(10), shape=(10,))\n",
    "print(x)\n",
    "y = jfunc_bis(x)\n",
    "y, jnp.sum(y)  # verify that it gives the same result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 115,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "SOH8JweLoGZZ",
    "outputId": "e89b1512-cfe5-4271-f172-d7fcdaf9678e"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 115,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.allclose(jax.jacfwd(jfunc_bis)(x),jacfwd_f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 117,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "wT0u-tfWp1FV",
    "outputId": "dfe03888-88d6-4bb1-edff-d8ccab88cd6c"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 117,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.allclose(jax.jacrev(jfunc_bis)(x),jacfwd_f)  # That's Ok Now !!!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "CCocRE2oJ636"
   },
   "source": [
    "# Takeaway message\n",
    "- Flux controls: use jax.lax (cond, while, fori_loop, scan)\n",
    "- Reminder that jax.numpy.where, jax.numpy.piecewise deal also to conditions over values in numpy arrays.\n",
    "- One can use auto-diff on function that uses jax.lax control flow\n",
    "- You have seen printing for debug of JIT functions\n",
    "- Sometimes you need to turn off the tracing of arguments for JIT\n",
    "- While is tricky: JIT or jacrev will be a problem ! impose to think differently\n",
    "- As thomb rule when you hesitate between fori_loop/while_loop or scan, then **\"always scan when you can!\"**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "djT19fVcqMmV"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "authorship_tag": "ABX9TyMB1BEEpi03Osict8Eo9RsZ",
   "include_colab_link": true,
   "provenance": [],
   "toc_visible": true
  },
  "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.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}