JAX_static_traced_var_func.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_static_traced_var_func.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": "uVx_yyP68jC9"
   },
   "outputs": [],
   "source": [
    "import jax\n",
    "from jax import lax, grad, jit, vmap, jacrev\n",
    "from jax.tree_util import register_pytree_node_class\n",
    "\n",
    "import jax.numpy as jnp\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "2Mr-9-fVDH7b"
   },
   "outputs": [],
   "source": [
    "import jaxlib.xla_extension as xla_ext\n",
    "\n",
    "# You could use the presets\n",
    "#option = xla_ext.HloPrintOptions.short_parsable()\n",
    "#option = xla_ext.HloPrintOptions.canonical()\n",
    "#option = xla_ext.HloPrintOptions.fingerprint()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "1OFoiZmk8-gK"
   },
   "source": [
    "# Topic: pure numpy & JAX\n",
    "\n",
    "The question arise sometimes after playing with JAX: does one can use pure numpy inside a JAX function? why sometimes one NEEDS to use pure numpy even in a Jitable JAX function... Lets us take some use-cases.\n",
    "\n",
    "One things that is stated in the JAX doc [Thinking in Jax](https://jax.readthedocs.io/en/latest/notebooks/thinking_in_jax.html) is related to JIT:   \n",
    "\n",
    "    \"Not all JAX code can be JIT compiled, as it requires array shapes to be static & known at compile time\"\n",
    "\n",
    "This is more in action with a JAX change done after version 0.2 (currently JAX version is >= 0.4.6).\n",
    "\n",
    "Key Concept:\n",
    "    \n",
    "    Just as values can be either static or traced, operations can be static or traced.\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Mo11RmDKHF5L"
   },
   "source": [
    "# A toy exemple"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "EHOzVNpppQh2"
   },
   "source": [
    "Let us consider this Numpy code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "YpVYXvxeGmF9",
    "outputId": "bb2a23c9-1036-4afa-803d-94f34777041f"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "in=  [[ 0.  1.  2.  3.]\n",
      " [ 4.  5.  6.  7.]\n",
      " [ 8.  9. 10. 11.]]\n",
      "out=  [ 0.  1.  2.  3.  4.  5.  6.  7.  8.  9. 10. 11.]  sum: 66.0\n"
     ]
    }
   ],
   "source": [
    "x = np.arange(0,12.,1).reshape((3, 4))\n",
    "print(\"in= \",x)\n",
    "x = x.reshape(np.prod(x.shape))\n",
    "print(\"out= \",x, \" sum:\", x.sum())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "aVV0Nf2_pUOt"
   },
   "source": [
    "It is then natural to consider this JAX translation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "shk1jD4gB3pu"
   },
   "outputs": [],
   "source": [
    "def ex1(x):\n",
    "    size = jnp.prod(jnp.array(x.shape))\n",
    "    print(\"size:\",size)\n",
    "    return jnp.sum(x.reshape((size,)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "NUULcCZdHa17",
    "outputId": "2d5368d5-aee7-4830-cfbb-4222666d68a3"
   },
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'jnp' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[3], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m x \u001b[38;5;241m=\u001b[39m \u001b[43mjnp\u001b[49m\u001b[38;5;241m.\u001b[39marange(\u001b[38;5;241m0\u001b[39m,\u001b[38;5;241m12.\u001b[39m,\u001b[38;5;241m1\u001b[39m)\u001b[38;5;241m.\u001b[39mreshape((\u001b[38;5;241m3\u001b[39m, \u001b[38;5;241m4\u001b[39m))\n\u001b[1;32m      2\u001b[0m ex1(x)\n",
      "\u001b[0;31mNameError\u001b[0m: name 'jnp' is not defined"
     ]
    }
   ],
   "source": [
    "x = jnp.arange(0,12.,1).reshape((3, 4))\n",
    "ex1(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "BPOtR33qHsLA"
   },
   "source": [
    "ok all good. Let jit this function....  :)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 505
    },
    "id": "ncwdW5tTGfie",
    "outputId": "5be28431-206a-46cc-84cd-4f9c81d14ca1"
   },
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "Shapes must be 1D sequences of concrete values of integer type, got (Traced<ShapedArray(int32[])>with<DynamicJaxprTrace(level=1/0)>,).\nIf using `jit`, try using `static_argnums` or applying `jit` to smaller subfunctions.\nThe error occurred while tracing the function ex1 at <ipython-input-4-e59efdb7a0a7>:1 for jit. This value became a tracer due to JAX operations on these lines:\n\n  operation a:i32[2] = convert_element_type[new_dtype=int32 weak_type=False] b\n    from line <ipython-input-4-e59efdb7a0a7>:2 (ex1)",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-6-5fa6b22786ec>\u001b[0m in \u001b[0;36m<cell line: 1>\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mjax\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mjit\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mex1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;31m# a crash my god!\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-4-e59efdb7a0a7>\u001b[0m in \u001b[0;36mex1\u001b[0;34m(x)\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mex1\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[1;32m      2\u001b[0m   \u001b[0msize\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mjnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mprod\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mjnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0marray\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\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[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 3\u001b[0;31m   \u001b[0;32mreturn\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[0mreshape\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0msize\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[0m",
      "    \u001b[0;31m[... skipping hidden 2 frame]\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_compute_newshape\u001b[0;34m(a, newshape)\u001b[0m\n\u001b[1;32m    114\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    115\u001b[0m     \u001b[0mnewshape\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mnewshape\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 116\u001b[0;31m   \u001b[0mnewshape\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcore\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcanonicalize_shape\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnewshape\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# type: ignore[arg-type]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    117\u001b[0m   \u001b[0mneg1s\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mi\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mi\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0md\u001b[0m \u001b[0;32min\u001b[0m \u001b[0menumerate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnewshape\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mtype\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0md\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0mint\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0md\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[0m\n\u001b[1;32m    118\u001b[0m   \u001b[0;32mif\u001b[0m \u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mneg1s\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[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/jax/_src/core.py\u001b[0m in \u001b[0;36mcanonicalize_shape\u001b[0;34m(shape, context)\u001b[0m\n\u001b[1;32m   2140\u001b[0m   \u001b[0;32mexcept\u001b[0m \u001b[0mTypeError\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2141\u001b[0m     \u001b[0;32mpass\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 2142\u001b[0;31m   \u001b[0;32mraise\u001b[0m \u001b[0m_invalid_shape_error\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mshape\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcontext\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   2143\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2144\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mcanonicalize_dim\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0md\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mDimSize\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcontext\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mstr\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m\"\"\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0mDimSize\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mTypeError\u001b[0m: Shapes must be 1D sequences of concrete values of integer type, got (Traced<ShapedArray(int32[])>with<DynamicJaxprTrace(level=1/0)>,).\nIf using `jit`, try using `static_argnums` or applying `jit` to smaller subfunctions.\nThe error occurred while tracing the function ex1 at <ipython-input-4-e59efdb7a0a7>:1 for jit. This value became a tracer due to JAX operations on these lines:\n\n  operation a:i32[2] = convert_element_type[new_dtype=int32 weak_type=False] b\n    from line <ipython-input-4-e59efdb7a0a7>:2 (ex1)"
     ]
    }
   ],
   "source": [
    "jax.jit(ex1)(x) # a crash my god!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5A5uk_HNH3Sq"
   },
   "source": [
    "There is a crash because of \"reshape\" as  \"***Shapes must be 1D sequences of concrete values of integer type***\". But the solution to put \"x\" as static is not our solution as we want for instance compute gradiant, vmap...\n",
    "\n",
    "In fact: x.shape is performed statically, but jnp.prod is traced as can be shown running the following code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "id": "woEiwnWIqCXe"
   },
   "outputs": [],
   "source": [
    "def ex1_debug(x):\n",
    "  size = jnp.prod(jnp.array(x.shape))\n",
    "  print(f\"x = {x}\")\n",
    "  print(f\"x.shape = {x.shape}\")   # not traced\n",
    "  print(f\"jnp.array(x.shape).prod() = {size}\")\n",
    "  return jnp.sum(x.reshape((size,)))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 560
    },
    "id": "VsWNV9PQqTS3",
    "outputId": "807b5382-2daf-42ad-b21e-3753ce7219e3"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x = Traced<ShapedArray(float32[3,4])>with<DynamicJaxprTrace(level=1/0)>\n",
      "x.shape = (3, 4)\n",
      "jnp.array(x.shape).prod() = Traced<ShapedArray(int32[])>with<DynamicJaxprTrace(level=1/0)>\n"
     ]
    },
    {
     "ename": "TypeError",
     "evalue": "Shapes must be 1D sequences of concrete values of integer type, got (Traced<ShapedArray(int32[])>with<DynamicJaxprTrace(level=1/0)>,).\nIf using `jit`, try using `static_argnums` or applying `jit` to smaller subfunctions.\nThe error occurred while tracing the function ex1_debug at <ipython-input-7-2fed1eae02b0>:1 for jit. This value became a tracer due to JAX operations on these lines:\n\n  operation a:i32[2] = convert_element_type[new_dtype=int32 weak_type=False] b\n    from line <ipython-input-7-2fed1eae02b0>:2 (ex1_debug)",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-8-aec215d994ba>\u001b[0m in \u001b[0;36m<cell line: 1>\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mjax\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mjit\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mex1_debug\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# still crash but see the prints\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-7-2fed1eae02b0>\u001b[0m in \u001b[0;36mex1_debug\u001b[0;34m(x)\u001b[0m\n\u001b[1;32m      4\u001b[0m   \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34mf\"x.shape = {x.shape}\"\u001b[0m\u001b[0;34m)\u001b[0m   \u001b[0;31m# not traced\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m   \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34mf\"jnp.array(x.shape).prod() = {size}\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 6\u001b[0;31m   \u001b[0;32mreturn\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[0mreshape\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0msize\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[0m",
      "    \u001b[0;31m[... skipping hidden 2 frame]\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_compute_newshape\u001b[0;34m(a, newshape)\u001b[0m\n\u001b[1;32m    114\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    115\u001b[0m     \u001b[0mnewshape\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mnewshape\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 116\u001b[0;31m   \u001b[0mnewshape\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcore\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcanonicalize_shape\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnewshape\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# type: ignore[arg-type]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    117\u001b[0m   \u001b[0mneg1s\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mi\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mi\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0md\u001b[0m \u001b[0;32min\u001b[0m \u001b[0menumerate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnewshape\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mtype\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0md\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0mint\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0md\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[0m\n\u001b[1;32m    118\u001b[0m   \u001b[0;32mif\u001b[0m \u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mneg1s\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[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/jax/_src/core.py\u001b[0m in \u001b[0;36mcanonicalize_shape\u001b[0;34m(shape, context)\u001b[0m\n\u001b[1;32m   2140\u001b[0m   \u001b[0;32mexcept\u001b[0m \u001b[0mTypeError\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2141\u001b[0m     \u001b[0;32mpass\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 2142\u001b[0;31m   \u001b[0;32mraise\u001b[0m \u001b[0m_invalid_shape_error\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mshape\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcontext\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   2143\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2144\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mcanonicalize_dim\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0md\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mDimSize\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcontext\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mstr\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m\"\"\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0mDimSize\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mTypeError\u001b[0m: Shapes must be 1D sequences of concrete values of integer type, got (Traced<ShapedArray(int32[])>with<DynamicJaxprTrace(level=1/0)>,).\nIf using `jit`, try using `static_argnums` or applying `jit` to smaller subfunctions.\nThe error occurred while tracing the function ex1_debug at <ipython-input-7-2fed1eae02b0>:1 for jit. This value became a tracer due to JAX operations on these lines:\n\n  operation a:i32[2] = convert_element_type[new_dtype=int32 weak_type=False] b\n    from line <ipython-input-7-2fed1eae02b0>:2 (ex1_debug)"
     ]
    }
   ],
   "source": [
    "jax.jit(ex1_debug)(x)  # still crash but see the prints"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "id": "r5k6NBj_G-YH"
   },
   "outputs": [],
   "source": [
    "@jit\n",
    "def ex1_sol(x):\n",
    "  size = np.prod(np.array(x.shape))     #<= notice pure numpy used\n",
    "  return jnp.sum(x.reshape((size,)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "6y0hJovsIXpL",
    "outputId": "e024e4b9-a1e5-4aec-ab0c-f70f816e1cb4"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(66., dtype=float32)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = jnp.arange(0,12.,1).reshape((3, 4))\n",
    "ex1_sol(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "pnfnA7U8Ic9c",
    "outputId": "e7d5a4f9-76be-46a2-b124-c7302fa35a36"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([[1., 1., 1., 1.],\n",
       "       [1., 1., 1., 1.],\n",
       "       [1., 1., 1., 1.]], dtype=float32)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "grad(ex1_sol)(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "ORLy0j69Sb1H",
    "outputId": "f27a75ab-865f-42d0-e3a8-99fceee47e4b"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([66.     , -6.51921], dtype=float32)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#keep that vmap should be applyied on same sized objects\n",
    "vmap(ex1_sol)(jnp.array([x,jax.random.normal(jax.random.PRNGKey(10), shape=(12,)).reshape((3, 4))]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-h6QSmiItkBy"
   },
   "source": [
    "# An other toy example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "id": "iFPkZhcmLlUN"
   },
   "outputs": [],
   "source": [
    "def exo2(x):\n",
    "  n = x.shape[0]        # this is static Ok\n",
    "  print(\"x: \",x, type(x))\n",
    "  jax.debug.print(\"debug x: {}\",x)\n",
    "  print(\"n: \",n)\n",
    "  y = jnp.arange(0.,n)   # so ok\n",
    "  return jnp.sum(x*y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "nXoSiDMKOGMo",
    "outputId": "3bde5b5c-d14f-4424-d3ca-303295d0445f"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x:  [ 0.4676388  -0.26684377  1.1685165  -0.48255268 -1.4658067   0.408608\n",
      " -0.12636082 -0.16064268  0.876103   -0.23290116] <class 'jaxlib.xla_extension.ArrayImpl'>\n",
      "debug x: [ 0.4676388  -0.26684377  1.1685165  -0.48255268 -1.4658067   0.408608\n",
      " -0.12636082 -0.16064268  0.876103   -0.23290116]\n",
      "n:  10\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Array(-0.16760588, dtype=float32)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = jax.random.normal(jax.random.PRNGKey(10), shape=(10,))\n",
    "(exo2)(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "peN3vTjhtZfi",
    "outputId": "79bb15eb-db00-4d8d-8f5a-9bf500ab3d88"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x:  Traced<ShapedArray(float32[10])>with<DynamicJaxprTrace(level=1/0)> <class 'jax._src.interpreters.partial_eval.DynamicJaxprTracer'>\n",
      "n:  10\n",
      "debug x: [ 0.4676388  -0.26684377  1.1685165  -0.48255268 -1.4658067   0.408608\n",
      " -0.12636082 -0.16064268  0.876103   -0.23290116]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Array(-0.16760588, dtype=float32)"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jax.jit(exo2)(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "538rzaoYtcsK",
    "outputId": "16a9a008-b41e-4a68-925c-ec9f99d19892"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x:  Traced<ShapedArray(float32[10])>with<JVPTrace(level=3/0)> with\n",
      "  primal = Traced<ShapedArray(float32[10])>with<DynamicJaxprTrace(level=1/0)>\n",
      "  tangent = Traced<ShapedArray(float32[10])>with<JaxprTrace(level=2/0)> with\n",
      "    pval = (ShapedArray(float32[10]), None)\n",
      "    recipe = LambdaBinding() <class 'jax._src.interpreters.ad.JVPTracer'>\n",
      "n:  10\n",
      "debug x: [ 0.4676388  -0.26684377  1.1685165  -0.48255268 -1.4658067   0.408608\n",
      " -0.12636082 -0.16064268  0.876103   -0.23290116]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Array([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.], dtype=float32)"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jax.jit(jax.grad(exo2))(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "id": "P2NqXAh8t3y9"
   },
   "outputs": [],
   "source": [
    "def exo2_modif(x,n):\n",
    "  print(\"x: \",x, type(x))\n",
    "  print(\"n: \",n)\n",
    "  y = jnp.arange(0.,n)   # No more ok as n is now traced\n",
    "  return jnp.sum(x*y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 468
    },
    "id": "Mrg4AyhduE9W",
    "outputId": "b2cd443e-8c22-4f61-9965-ae49da26aabd"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x:  Traced<ShapedArray(float32[10])>with<DynamicJaxprTrace(level=1/0)> <class 'jax._src.interpreters.partial_eval.DynamicJaxprTracer'>\n",
      "n:  Traced<ShapedArray(int32[], weak_type=True)>with<DynamicJaxprTrace(level=1/0)>\n"
     ]
    },
    {
     "ename": "ConcretizationTypeError",
     "evalue": "Abstract tracer value encountered where concrete value is expected: traced array with shape int32[].\nIt arose in the jnp.arange argument 'stop'\nThe error occurred while tracing the function exo2_modif at <ipython-input-17-10dfe816faa6>:1 for jit. This concrete value was not available in Python because it depends on the value of the argument n.\n\nSee https://jax.readthedocs.io/en/latest/errors.html#jax.errors.ConcretizationTypeError",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mConcretizationTypeError\u001b[0m                   Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-18-a08c46bac412>\u001b[0m in \u001b[0;36m<cell line: 1>\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mjax\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mjit\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mexo2_modif\u001b[0m\u001b[0;34m)\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[0mshape\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\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;31m[... skipping hidden 11 frame]\u001b[0m\n",
      "\u001b[0;32m<ipython-input-17-10dfe816faa6>\u001b[0m in \u001b[0;36mexo2_modif\u001b[0;34m(x, n)\u001b[0m\n\u001b[1;32m      2\u001b[0m   \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"x: \"\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtype\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[1;32m      3\u001b[0m   \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"n: \"\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[0m\n\u001b[0;32m----> 4\u001b[0;31m   \u001b[0my\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mjnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0marange\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[0;31m# No more ok as n is now traced\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      5\u001b[0m   \u001b[0;32mreturn\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[0my\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/lax_numpy.py\u001b[0m in \u001b[0;36marange\u001b[0;34m(start, stop, step, dtype)\u001b[0m\n\u001b[1;32m   2500\u001b[0m       \u001b[0mstart\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcore\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mconcrete_or_error\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstart\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m\"It arose in the jnp.arange argument 'start'\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2501\u001b[0m   \u001b[0mutil\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcheck_arraylike_or_none\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"arange\"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstop\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstep\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 2502\u001b[0;31m   \u001b[0mstop\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcore\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mconcrete_or_error\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstop\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m\"It arose in the jnp.arange argument 'stop'\"\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   2503\u001b[0m   \u001b[0mstep\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcore\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mconcrete_or_error\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstep\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m\"It arose in the jnp.arange argument 'step'\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2504\u001b[0m   \u001b[0mstart_name\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m\"stop\"\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mstop\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mNone\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0mstep\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mNone\u001b[0m \u001b[0;32melse\u001b[0m \u001b[0;34m\"start\"\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/core.py\u001b[0m in \u001b[0;36mconcrete_or_error\u001b[0;34m(force, val, context)\u001b[0m\n\u001b[1;32m   1507\u001b[0m       \u001b[0;32mreturn\u001b[0m \u001b[0mforce\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mval\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0maval\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mval\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1508\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-> 1509\u001b[0;31m       \u001b[0;32mraise\u001b[0m \u001b[0mConcretizationTypeError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mval\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcontext\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   1510\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   1511\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0mforce\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mval\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mConcretizationTypeError\u001b[0m: Abstract tracer value encountered where concrete value is expected: traced array with shape int32[].\nIt arose in the jnp.arange argument 'stop'\nThe error occurred while tracing the function exo2_modif at <ipython-input-17-10dfe816faa6>:1 for jit. This concrete value was not available in Python because it depends on the value of the argument n.\n\nSee https://jax.readthedocs.io/en/latest/errors.html#jax.errors.ConcretizationTypeError"
     ]
    }
   ],
   "source": [
    "jax.jit(exo2_modif)(x, x.shape[0])  # Crash !!!!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "id": "gRn4AEwquaLY"
   },
   "outputs": [],
   "source": [
    "def exo2_modif(x,n):\n",
    "  print(\"x: \",x, type(x))\n",
    "  print(\"n: \",n)\n",
    "  y = jnp.arange(0.,n)\n",
    "  return jnp.sum(x*y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "z6T8jrqRuzHl",
    "outputId": "30b2bfea-0828-4b85-95d4-d817a46a5bbc"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x:  Traced<ShapedArray(float32[10])>with<DynamicJaxprTrace(level=1/0)> <class 'jax._src.interpreters.partial_eval.DynamicJaxprTracer'>\n",
      "n:  10\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Array(-0.16760588, dtype=float32)"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jax.jit(exo2_modif, static_argnums=(1,))(x, x.shape[0]) # Now n is no more traced by construction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "ipH3a8nQvLcx",
    "outputId": "b5718d42-f4c4-4d2b-cec7-56dd06eca268"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x:  Traced<ShapedArray(float32[10])>with<JVPTrace(level=3/0)> with\n",
      "  primal = Traced<ShapedArray(float32[10])>with<DynamicJaxprTrace(level=1/0)>\n",
      "  tangent = Traced<ShapedArray(float32[10])>with<JaxprTrace(level=2/0)> with\n",
      "    pval = (ShapedArray(float32[10]), None)\n",
      "    recipe = LambdaBinding() <class 'jax._src.interpreters.ad.JVPTracer'>\n",
      "n:  10\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Array([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.], dtype=float32)"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jax.jit(jax.grad(exo2_modif), static_argnums=(1,))(x, x.shape[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Hboq3skQ0TpW"
   },
   "source": [
    "# More advanced exo\n",
    "\n",
    "The above exercises could be concidered as toy ex. but the kind of pb can be the cause of crash not so eveident to debug if not aware of.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "id": "QChi0nsO0X5_"
   },
   "outputs": [],
   "source": [
    "def ex_convolve(in1,in2):\n",
    "\n",
    "  axes = range(in1.ndim)\n",
    "\n",
    "  s1 = in1.shape\n",
    "  s2 = in2.shape\n",
    "  shape = [s1[i] + s2[i] -1 for i in range(in1.ndim)]\n",
    "\n",
    "  print(\"shape\",shape)\n",
    "\n",
    "  fshape = 2**jnp.ceil(jnp.log2(jnp.array(shape))).astype('uint16')  #<-------\n",
    "\n",
    "  print(\"fshape\",fshape)\n",
    "\n",
    "  sp1 = jnp.fft.fftn(in1, fshape, axes=axes)                         #<-------\n",
    "  sp2 = jnp.fft.fftn(in2, fshape, axes=axes)\n",
    "  conv = jnp.fft.irfftn(sp1 * sp2, fshape, axes=axes)\n",
    "  conv = conv[tuple(map(slice, shape))]\n",
    "\n",
    "  def _centered(arr, newshape):\n",
    "    # Return the center newshape portion of the array.\n",
    "    newshape = np.asarray(newshape)\n",
    "    currshape = np.array(arr.shape)\n",
    "    startind = (currshape - newshape) // 2\n",
    "    endind = startind + newshape\n",
    "    myslice = [slice(startind[k], endind[k]) for k in range(len(endind))]\n",
    "    return arr[tuple(myslice)]\n",
    "\n",
    "  return _centered(conv, s1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Y8gnG5PL0kUR",
    "outputId": "ceb82b96-c058-4dc1-e225-cb2abb8429ff"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "shape [19, 19]\n",
      "fshape [32 32]\n"
     ]
    }
   ],
   "source": [
    "key, key1, key2= jax.random.split(jax.random.PRNGKey(1),3)\n",
    "x1 = jax.random.normal(key1, shape=(10,10))\n",
    "x2 = jax.random.normal(key2, shape=(10,10))\n",
    "res = ex_convolve(x1,x2)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 782
    },
    "id": "x6nmCmj82k-F",
    "outputId": "b1796332-62ce-489b-d025-7b4412708a8a"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "shape [19, 19]\n",
      "fshape Traced<ShapedArray(uint16[2])>with<DynamicJaxprTrace(level=1/0)>\n"
     ]
    },
    {
     "ename": "TracerIntegerConversionError",
     "evalue": "The __index__() method was called on traced array with shape uint16[].\nThe error occurred while tracing the function ex_convolve at <ipython-input-22-47be1225c5fd>:1 for jit. This value became a tracer due to JAX operations on these lines:\n\n  operation a:i32[2] = convert_element_type[new_dtype=int32 weak_type=False] b\n    from line <ipython-input-22-47be1225c5fd>:11 (ex_convolve)\n\n  operation a:f32[] = log b\n    from line <ipython-input-22-47be1225c5fd>:11 (ex_convolve)\n\n  operation a:u16[] = convert_element_type[new_dtype=uint16 weak_type=False] b\n    from line <ipython-input-22-47be1225c5fd>:11 (ex_convolve)\n\n  operation a:bool[] = eq b c\n    from line <ipython-input-22-47be1225c5fd>:11 (ex_convolve)\n\n  operation a:bool[2] = ne b c\n    from line <ipython-input-22-47be1225c5fd>:11 (ex_convolve)\n\n(Additional originating lines are not shown.)\nSee https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerIntegerConversionError",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTracerIntegerConversionError\u001b[0m              Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-24-256d6996102b>\u001b[0m in \u001b[0;36m<cell line: 1>\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mres1\u001b[0m\u001b[0;34m=\u001b[0m \u001b[0mjax\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mjit\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mex_convolve\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx1\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mx2\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;31m[... skipping hidden 11 frame]\u001b[0m\n",
      "\u001b[0;32m<ipython-input-22-47be1225c5fd>\u001b[0m in \u001b[0;36mex_convolve\u001b[0;34m(in1, in2)\u001b[0m\n\u001b[1;32m     13\u001b[0m   \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"fshape\"\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mfshape\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     14\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 15\u001b[0;31m   \u001b[0msp1\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mjnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfft\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfftn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0min1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfshape\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0maxes\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0maxes\u001b[0m\u001b[0;34m)\u001b[0m                         \u001b[0;31m#<-------\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     16\u001b[0m   \u001b[0msp2\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mjnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfft\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfftn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0min2\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfshape\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0maxes\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0maxes\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     17\u001b[0m   \u001b[0mconv\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mjnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfft\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mirfftn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0msp1\u001b[0m \u001b[0;34m*\u001b[0m \u001b[0msp2\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfshape\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0maxes\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0maxes\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/fft.py\u001b[0m in \u001b[0;36mfftn\u001b[0;34m(a, s, axes, norm)\u001b[0m\n\u001b[1;32m    114\u001b[0m          \u001b[0maxes\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mSequence\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mint\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m|\u001b[0m \u001b[0;32mNone\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[1;32m    115\u001b[0m          norm: str | None = None) -> Array:\n\u001b[0;32m--> 116\u001b[0;31m   \u001b[0;32mreturn\u001b[0m \u001b[0m_fft_core\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'fftn'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mxla_client\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mFftType\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mFFT\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0ma\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0ms\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0maxes\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnorm\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    117\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    118\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/jax/_src/numpy/fft.py\u001b[0m in \u001b[0;36m_fft_core\u001b[0;34m(func_name, fft_type, a, s, axes, norm)\u001b[0m\n\u001b[1;32m     53\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     54\u001b[0m   \u001b[0;32mif\u001b[0m \u001b[0ms\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---> 55\u001b[0;31m     \u001b[0ms\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtuple\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmap\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[0ms\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     56\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0many\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mless\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0ms\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m0\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     57\u001b[0m       \u001b[0;32mraise\u001b[0m \u001b[0mValueError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"Shape should be non-negative.\"\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   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;32m-> 1495\u001b[0;31m       \u001b[0;32mraise\u001b[0m \u001b[0mTracerIntegerConversionError\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   1496\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   1497\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;31mTracerIntegerConversionError\u001b[0m: The __index__() method was called on traced array with shape uint16[].\nThe error occurred while tracing the function ex_convolve at <ipython-input-22-47be1225c5fd>:1 for jit. This value became a tracer due to JAX operations on these lines:\n\n  operation a:i32[2] = convert_element_type[new_dtype=int32 weak_type=False] b\n    from line <ipython-input-22-47be1225c5fd>:11 (ex_convolve)\n\n  operation a:f32[] = log b\n    from line <ipython-input-22-47be1225c5fd>:11 (ex_convolve)\n\n  operation a:u16[] = convert_element_type[new_dtype=uint16 weak_type=False] b\n    from line <ipython-input-22-47be1225c5fd>:11 (ex_convolve)\n\n  operation a:bool[] = eq b c\n    from line <ipython-input-22-47be1225c5fd>:11 (ex_convolve)\n\n  operation a:bool[2] = ne b c\n    from line <ipython-input-22-47be1225c5fd>:11 (ex_convolve)\n\n(Additional originating lines are not shown.)\nSee https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerIntegerConversionError"
     ]
    }
   ],
   "source": [
    "res1= jax.jit(ex_convolve)(x1,x2) # crash"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "QVVNcCOd3TMm"
   },
   "source": [
    "The problem comes from the fshape used by fft is Traced while static is required as \"Shape = Sequence[int]\". The solution is to use pure Numpy\n",
    "function to compute fshape."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "id": "DPA8BjVw6gNc"
   },
   "outputs": [],
   "source": [
    "def ex_convolve_sol(in1,in2):\n",
    "\n",
    "  axes = range(in1.ndim)\n",
    "\n",
    "  s1 = in1.shape\n",
    "  s2 = in2.shape\n",
    "  shape = [s1[i] + s2[i] -1 for i in range(in1.ndim)]\n",
    "\n",
    "  print(\"shape\",shape)\n",
    "\n",
    "  fshape = 2**np.ceil(np.log2(np.array(shape))).astype('uint16') # ok now\n",
    "\n",
    "  print(\"fshape\",fshape)\n",
    "\n",
    "  sp1 = jnp.fft.fftn(in1, fshape, axes=axes)\n",
    "  sp2 = jnp.fft.fftn(in2, fshape, axes=axes)\n",
    "  conv = jnp.fft.irfftn(sp1 * sp2, fshape, axes=axes)\n",
    "  conv = conv[tuple(map(slice, shape))]\n",
    "\n",
    "  def _centered(arr, newshape):\n",
    "    # Return the center newshape portion of the array.\n",
    "    newshape = np.asarray(newshape)\n",
    "    currshape = np.array(arr.shape)\n",
    "    startind = (currshape - newshape) // 2\n",
    "    endind = startind + newshape\n",
    "    myslice = [slice(startind[k], endind[k]) for k in range(len(endind))]\n",
    "    return arr[tuple(myslice)]\n",
    "\n",
    "  return _centered(conv, s1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "LguUoKKG6pYy",
    "outputId": "fcd1ceb0-5c36-492b-e105-188c0bf3cc18"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "shape [19, 19]\n",
      "fshape [32 32]\n"
     ]
    }
   ],
   "source": [
    "res1= jax.jit(ex_convolve_sol)(x1,x2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Ap_dShAXUCfL"
   },
   "source": [
    "# Takeaway message\n",
    "\n",
    "- Properties of arrays like shapes, sizes, and dtypes are always static within JAX, so they can be used to define other static quantities\n",
    "\n",
    "- Not all JAX code can be JIT compiled, as it requires array shapes to be static & known at compile time.\n",
    "\n",
    "- Static operations are evaluated at compile-time in Python; traced operations are compiled & evaluated at run-time in XLA.\n",
    "\n",
    "- Use numpy for operations that you want to be static; use jax.numpy for operations that you want to be traced.\n",
    "\n",
    "\n",
    "- This post was in fact the result of an exchange I had with JAX developpers (end March 2023). The [JAX doc \"how to think in JAX\"] (https://jax.readthedocs.io/en/latest/notebooks/thinking_in_jax.html#how-to-think-in-jax) was completed after this post. As the JAX API may change, this doc can also change too.\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "uvlO8GJCUFI-"
   },
   "source": [
    "# Epilog"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "7QjYSMEzamX7",
    "outputId": "edb7c757-58c7-4f13-e4a1-731434b80ad3"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x: Traced<ShapedArray(float32[3,4])>with<DynamicJaxprTrace(level=1/0)>\n",
      "66.0\n",
      "x: Traced<ConcreteArray([[ 0.  1.  2.  3.]\n",
      " [ 4.  5.  6.  7.]\n",
      " [ 8.  9. 10. 11.]], dtype=float32)>with<JVPTrace(level=2/0)> with\n",
      "  primal = array([[ 0.,  1.,  2.,  3.],\n",
      "       [ 4.,  5.,  6.,  7.],\n",
      "       [ 8.,  9., 10., 11.]])\n",
      "  tangent = Traced<ShapedArray(float32[3,4])>with<JaxprTrace(level=1/0)> with\n",
      "    pval = (ShapedArray(float32[3,4]), None)\n",
      "    recipe = LambdaBinding()\n",
      "[[1. 1. 1. 1.]\n",
      " [1. 1. 1. 1.]\n",
      " [1. 1. 1. 1.]]\n"
     ]
    }
   ],
   "source": [
    "# a curiosity due to np.sum beahvior\n",
    "def ex0(x):\n",
    "  print(\"x:\",x)\n",
    "  size = np.prod(np.array(x.shape))\n",
    "  return np.sum(x.reshape((size,)))\n",
    "print(jit(ex0)(np.arange(0,12.,1).reshape((3, 4))))\n",
    "print(grad(ex0)(np.arange(0,12.,1).reshape((3, 4))))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "mSXWh1i_V62l"
   },
   "source": [
    "A you probably noticed the `ex0` code rely only on pure Numpy function, so one can ask how is it possible to trace `x` and make possible the JIT and Autodiff???\n",
    "\n",
    "The reason is very tricky: `np.sum` looks in fact if its argument has a `.sum()` method, if so it calls it. And JAX tracers do have a sum method, so it is why every thing works ! now if you try np.exp on top of np.sum then there will be a crash... Theris no miracle.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "9vXpzVPNm-n3"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "authorship_tag": "ABX9TyOGtaWgKt51KnPnMLFcPjZd",
   "collapsed_sections": [
    "1OFoiZmk8-gK"
   ],
   "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
}