JAX_JIT_in_class.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_JIT_in_class.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": "ieUbBxWOdMxq"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import scipy as sc\n",
    "\n",
    "\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import jax.scipy as jsc\n",
    "\n",
    "from jax import grad, jit, vmap\n",
    "from jax import jacfwd, jacrev, hessian\n",
    "#from jax.ops import index, index_update\n",
    "jax.config.update(\"jax_enable_x64\", True)\n",
    "\n",
    "from functools import partial\n",
    "from typing import Union, Dict, Callable, Optional, Tuple, Any"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 36
    },
    "id": "I5OJV3UUfSm0",
    "outputId": "7b3186bf-0b34-4829-d1cd-3acd6e5324ed"
   },
   "outputs": [
    {
     "data": {
      "application/vnd.google.colaboratory.intrinsic+json": {
       "type": "string"
      },
      "text/plain": [
       "'0.4.26'"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jax.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "id": "wSf3EGXIewrV"
   },
   "outputs": [],
   "source": [
    "def u_print(idx:int, *args)->int:\n",
    "    print(f\"{idx}):\",*args)\n",
    "    idx +=1\n",
    "    return idx\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "l6OOn8bneA1v"
   },
   "source": [
    "# Topic: JIT and Class methods\n",
    "\n",
    "This notebook is a bit technical but deals with a real problem when one comes from the Object Oriented world (ex. C++ & Python).\n",
    "\n",
    "it's an element of a JAX thread that I initiated and that becomes a part of the doc [see FAQ](https://jax.readthedocs.io/en/latest/faq.html#strategy-1-jit-compiled-helper-function)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "NrV88f8gdS3K",
    "outputId": "bdd0dcaf-6ba0-46b4-adb7-1d447657472d"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "New A\n",
      "0): 2.0 None\n",
      "1): 2.0 20.0\n",
      "2): 2.0 22.0\n",
      "New A\n",
      "3): 3.0 None\n",
      "4): 3.0 60.0\n"
     ]
    }
   ],
   "source": [
    "class A():\n",
    "    def __init__(self, a: float):\n",
    "        print(\"New A\")\n",
    "        self.a = a                   # a variable that will be \"static\" once the obj is created\n",
    "        self.b = None                # a variable (\"dynamic\") updated by f\n",
    "\n",
    "    def f(self, var: float) -> None:\n",
    "        self.b = self.a * var\n",
    "\n",
    "######\n",
    "\n",
    "idp=0  # (just an index for printing)\n",
    "\n",
    "objA = A(2.0)\n",
    "idp = u_print(idp,objA.a, objA.b)\n",
    "objA.f(10.)\n",
    "idp = u_print(idp,objA.a, objA.b)\n",
    "objA.f(11.)\n",
    "idp = u_print(idp,objA.a, objA.b)\n",
    "\n",
    "\n",
    "objA = A(3.0)\n",
    "idp = u_print(idp,objA.a, objA.b)\n",
    "objA.f(20.)\n",
    "idp = u_print(idp,objA.a, objA.b)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7mEvPJ-Fo1-Q"
   },
   "source": [
    "# And now use jit on f"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 370
    },
    "id": "OTZ0gec4oVCw",
    "outputId": "dc34d866-f183-4105-9d63-a9376b6906c7"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0): 2.0 None\n"
     ]
    },
    {
     "ename": "TypeError",
     "evalue": "Cannot interpret value of type <class '__main__.A'> as an abstract array; it does not have a dtype attribute",
     "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-eaad4a8b0855>\u001b[0m in \u001b[0;36m<cell line: 15>\u001b[0;34m()\u001b[0m\n\u001b[1;32m     13\u001b[0m \u001b[0mobjA\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mA\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m2.0\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[0midp\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mu_print\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0midp\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mobjA\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0ma\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mobjA\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mb\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 15\u001b[0;31m \u001b[0mobjA\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mf\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 5 frame]\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/jax/_src/api_util.py\u001b[0m in \u001b[0;36m_shaped_abstractify_slow\u001b[0;34m(x)\u001b[0m\n\u001b[1;32m    586\u001b[0m     \u001b[0mdtype\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mdtypes\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcanonicalize_dtype\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdtype\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mallow_extended_dtype\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    587\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--> 588\u001b[0;31m     raise TypeError(\n\u001b[0m\u001b[1;32m    589\u001b[0m         \u001b[0;34mf\"Cannot interpret value of type {type(x)} as an abstract array; it \"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    590\u001b[0m         \"does not have a dtype attribute\")\n",
      "\u001b[0;31mTypeError\u001b[0m: Cannot interpret value of type <class '__main__.A'> as an abstract array; it does not have a dtype attribute"
     ]
    }
   ],
   "source": [
    "class A():\n",
    "    def __init__(self, a: float):\n",
    "        self.a = a                   # a variable that will be \"static\" once the obj is created\n",
    "        self.b = None                # a variable (\"dynamic\") updated by f\n",
    "\n",
    "    @jit\n",
    "    def f(self, var: float) -> None:\n",
    "        self.b = self.a * var\n",
    "\n",
    "\n",
    "idp=0\n",
    "\n",
    "objA = A(2.0)\n",
    "idp = u_print(idp,objA.a, objA.b)\n",
    "objA.f(10.)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "AnJM_GmTpP85"
   },
   "source": [
    "The problem is that the first argument of the function is **self**, which has the type `A`and JAX does not know what it is.\n",
    "\n",
    "There are different basic strategies that we can use in this case, and we will discuss them below."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4PHx5v9Lpo3C"
   },
   "source": [
    "# 1st strategy: external function ('helper')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "XeYg2G1MqYcx",
    "outputId": "d672e9ce-2ded-44f0-a677-cbca6d2a1072"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "New Obj A (1)\n",
      "0): 2.0 None\n",
      "compile...\n",
      "1): 2.0 20.0\n",
      "2): 2.0 22.0\n",
      "New Obj A (2)\n",
      "3): 3.0 None\n",
      "4): 3.0 60.0\n",
      "5): 3.0 63.0\n",
      "New Obj A (3)\n",
      "6): 4.0 None\n",
      "7): 4.0 80.0\n",
      "f is called with an array not a float => compile\n",
      "compile...\n",
      "8): 4.0 [80.]\n",
      "9): 4.0 [84.]\n",
      "10): 4.0 84.0\n",
      "change internal objA.a\n",
      "11): 400.0 8400.0\n"
     ]
    }
   ],
   "source": [
    "jax.clear_caches()  # clear all compilations\n",
    "\n",
    "@jit\n",
    "def _f(a, var):\n",
    "    print(\"compile...\")\n",
    "    res = a * var\n",
    "    return res\n",
    "\n",
    "class A():\n",
    "    def __init__(self, a):\n",
    "        self.a = a\n",
    "        self.b = None\n",
    "\n",
    "    def f(self, var):\n",
    "        self.b = _f(self.a,var)\n",
    "\n",
    "\n",
    "\n",
    "idp=0\n",
    "\n",
    "print(\"New Obj A (1)\")\n",
    "objA = A(2.0)\n",
    "idp = u_print(idp, objA.a, objA.b)\n",
    "objA.f(10.0)\n",
    "idp = u_print(idp, objA.a, objA.b)\n",
    "objA.f(11.0)\n",
    "idp = u_print(idp, objA.a, objA.b)\n",
    "\n",
    "\n",
    "print(\"New Obj A (2)\")\n",
    "objA = A(3.0)\n",
    "idp = u_print(idp, objA.a, objA.b)\n",
    "objA.f(20.)\n",
    "idp = u_print(idp, objA.a, objA.b)\n",
    "objA.f(21.)\n",
    "idp = u_print(idp, objA.a, objA.b)\n",
    "\n",
    "\n",
    "print(\"New Obj A (3)\")\n",
    "objA = A(4.0)\n",
    "idp = u_print(idp, objA.a, objA.b)\n",
    "objA.f(20.)\n",
    "idp = u_print(idp, objA.a, objA.b)\n",
    "\n",
    "print(\"f is called with an array not a float => compile\")\n",
    "objA.f(jnp.array([20.]))\n",
    "idp = u_print(idp, objA.a, objA.b)\n",
    "\n",
    "objA.f(jnp.array([21.]))\n",
    "idp = u_print(idp, objA.a, objA.b)\n",
    "\n",
    "objA.f(21.)\n",
    "idp = u_print(idp, objA.a, objA.b)\n",
    "\n",
    "print(\"change internal objA.a\")\n",
    "\n",
    "objA.a = 400.  # change\n",
    "objA.f(21.)\n",
    "idp = u_print(idp, objA.a, objA.b)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0mMrL3MdsuAj"
   },
   "source": [
    "## Review of the helper method\n",
    "\n",
    "- This is a simple method\n",
    "implementation, and we don't have to instruct JAX how to use the A class.\n",
    "\n",
    "- Now it becomes a matter of taste to code a helper per function to use jit. But we can do file encapsulation so that at least the code of \"A\" is in the code of the definition of \"A\".\n",
    "\n",
    "- Pb: how to manage teh use-case with the function f needing for example another function (either member of A, or external to A)?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6qkJjt-JtSXg"
   },
   "source": [
    "# 2nd strategy: self as static\n",
    "\n",
    "\n",
    "it is a classic procedure that is often proposed in forum.\n",
    "But..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 800
    },
    "id": "wz3kq-Oqq5NF",
    "outputId": "599c1759-eff0-40f7-a7e1-3d554beca920"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "new A\n",
      "0): 2.0 None\n",
      "compile...\n",
      "1): 2.0 Traced<ShapedArray(float64[], weak_type=True)>with<DynamicJaxprTrace(level=1/0)>\n",
      "g...: Traced<ShapedArray(float64[], weak_type=True)>with<DynamicJaxprTrace(level=1/0)>\n"
     ]
    },
    {
     "ename": "UnexpectedTracerError",
     "evalue": "Encountered an unexpected tracer. A function transformed by JAX had a side effect, allowing for a reference to an intermediate value with type float64[] wrapped in a DynamicJaxprTracer to escape the scope of the transformation.\nJAX transformations require that functions explicitly return their outputs, and disallow saving intermediate values to global state.\nThe function being traced when the value leaked was f at <ipython-input-20-c8714ebe4d23>:9 traced for jit.\n------------------------------\nThe leaked intermediate value was created on line <ipython-input-20-c8714ebe4d23>:12 (f). \n------------------------------\nWhen the value was created, the final 5 stack frames (most recent last) excluding JAX-internal frames were:\n------------------------------\n/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py:3257 (run_cell_async)\n/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py:3473 (run_ast_nodes)\n/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py:3553 (run_code)\n<ipython-input-20-c8714ebe4d23>:22 (<cell line: 22>)\n<ipython-input-20-c8714ebe4d23>:12 (f)\n------------------------------\n\nTo catch the leak earlier, try setting the environment variable JAX_CHECK_TRACER_LEAKS or using the `jax.checking_leaks` context manager.\nSee https://jax.readthedocs.io/en/latest/errors.html#jax.errors.UnexpectedTracerError",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mUnexpectedTracerError\u001b[0m                     Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-20-c8714ebe4d23>\u001b[0m in \u001b[0;36m<cell line: 24>\u001b[0;34m()\u001b[0m\n\u001b[1;32m     22\u001b[0m \u001b[0mobjA\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mf\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[1;32m     23\u001b[0m \u001b[0midp\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mu_print\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0midp\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mobjA\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0ma\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mobjA\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mb\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 24\u001b[0;31m \u001b[0mres\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mobjA\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mg\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m   \u001b[0;31m# will case a crash but this is not the culpit\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     25\u001b[0m \u001b[0mprint\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[0m\n",
      "\u001b[0;32m<ipython-input-20-c8714ebe4d23>\u001b[0m in \u001b[0;36mg\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m     14\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mg\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\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     15\u001b[0m         \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"g...:\"\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mb\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 16\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mb\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mb\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     17\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     18\u001b[0m \u001b[0midp\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m0\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;36mdeferring_binary_op\u001b[0;34m(self, other)\u001b[0m\n\u001b[1;32m    262\u001b[0m     \u001b[0margs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mother\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mswap\u001b[0m \u001b[0;32melse\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mother\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    263\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mother\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0m_accepted_binop_types\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--> 264\u001b[0;31m       \u001b[0;32mreturn\u001b[0m \u001b[0mbinary_op\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    265\u001b[0m     \u001b[0;31m# Note: don't use isinstance here, because we don't want to raise for\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    266\u001b[0m     \u001b[0;31m# subclasses, e.g. NamedTuple objects that may override operators.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "    \u001b[0;31m[... skipping hidden 5 frame]\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/jax/_src/interpreters/partial_eval.py\u001b[0m in \u001b[0;36m_assert_live\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m   1718\u001b[0m   \u001b[0;32mdef\u001b[0m \u001b[0m_assert_live\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[1;32m   1719\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_trace\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmain\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mjaxpr_stack\u001b[0m\u001b[0;34m:\u001b[0m  \u001b[0;31m# type: ignore\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1720\u001b[0;31m       \u001b[0;32mraise\u001b[0m \u001b[0mcore\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mescaped_tracer_error\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\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   1721\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1722\u001b[0m   \u001b[0;32mdef\u001b[0m \u001b[0mget_referent\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\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;31mUnexpectedTracerError\u001b[0m: Encountered an unexpected tracer. A function transformed by JAX had a side effect, allowing for a reference to an intermediate value with type float64[] wrapped in a DynamicJaxprTracer to escape the scope of the transformation.\nJAX transformations require that functions explicitly return their outputs, and disallow saving intermediate values to global state.\nThe function being traced when the value leaked was f at <ipython-input-20-c8714ebe4d23>:9 traced for jit.\n------------------------------\nThe leaked intermediate value was created on line <ipython-input-20-c8714ebe4d23>:12 (f). \n------------------------------\nWhen the value was created, the final 5 stack frames (most recent last) excluding JAX-internal frames were:\n------------------------------\n/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py:3257 (run_cell_async)\n/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py:3473 (run_ast_nodes)\n/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py:3553 (run_code)\n<ipython-input-20-c8714ebe4d23>:22 (<cell line: 22>)\n<ipython-input-20-c8714ebe4d23>:12 (f)\n------------------------------\n\nTo catch the leak earlier, try setting the environment variable JAX_CHECK_TRACER_LEAKS or using the `jax.checking_leaks` context manager.\nSee https://jax.readthedocs.io/en/latest/errors.html#jax.errors.UnexpectedTracerError"
     ]
    }
   ],
   "source": [
    "jax.clear_caches()  # clear all compilations\n",
    "\n",
    "class A():\n",
    "    def __init__(self, a: float):\n",
    "        print(\"new A\")\n",
    "        self.a = a\n",
    "        self.b = None\n",
    "\n",
    "    @partial(jit, static_argnums=(0,))   # \"self\" is static\n",
    "    def f(self, var: float) -> None:\n",
    "        print(\"compile...\")\n",
    "        self.b = self.a * var  #<- this is the culpit of the later crash on \"g\"\n",
    "\n",
    "    def g(self):\n",
    "        print(\"g...:\",self.b)\n",
    "        return self.b * self.b\n",
    "\n",
    "idp=0\n",
    "\n",
    "objA = A(2.0)\n",
    "idp = u_print(idp, objA.a, objA.b)\n",
    "objA.f(10.)\n",
    "idp = u_print(idp, objA.a, objA.b)\n",
    "res = objA.g()   # will case a crash but this is not the culpit\n",
    "print(res)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "JojVQYrounlb"
   },
   "source": [
    "The code crash and seems to point the `g` function BUT in fact the problem comes from the jitted `f` function with self as static: it uses Traced<ShapedArray> to analyze the code, and propagate...\n",
    "Nb. the trace dump may change between JAX version...\n",
    "\n",
    "So\n",
    "```\n",
    "Do not use `self.var =` in a jitted function with self as static.\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "iaHLhqw_tm4a",
    "outputId": "53eb4b4c-c930-447b-c3f8-7968b6d2261a"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "New obj A\n",
      "0): 2.0 None\n",
      "compile...\n",
      "1): 2.0 20.0\n",
      "g...: 20.0\n",
      "400.0\n"
     ]
    }
   ],
   "source": [
    "class A():\n",
    "    def __init__(self, a):\n",
    "        print(\"New obj A\")\n",
    "        self.a = a\n",
    "        self.b = None\n",
    "\n",
    "    def set_a(self,x):\n",
    "        self.a = x\n",
    "        print(\"new a:\",self.a)\n",
    "\n",
    "    def set_b(self,x):                  # (*) call from extern\n",
    "        self.b = x\n",
    "\n",
    "    @partial(jit, static_argnums=(0,))   #  \"self\" is static\n",
    "    def f(self, var):\n",
    "        print(\"compile...\")\n",
    "        return self.a * var             # see (*)\n",
    "\n",
    "    def g(self):\n",
    "        print(\"g...:\",self.b)\n",
    "        return self.b * self.b\n",
    "\n",
    "\n",
    "idp=0\n",
    "\n",
    "objA = A(2.0)\n",
    "idp = u_print(idp, objA.a, objA.b)\n",
    "b  = objA.f(10.)\n",
    "objA.set_b(b)      #<----------------------\n",
    "idp = u_print(idp, objA.a, objA.b)\n",
    "res = objA.g()\n",
    "print(res)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "PjTeWoM4Rbtl",
    "outputId": "fa87a97e-81aa-4b24-d3aa-4e070defcc47"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "new a: 100.0\n",
      "5): 100.0 20.0\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Array(20., dtype=float64, weak_type=True)"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "objA.set_a(100.)   # <------ change \"a\"\n",
    "idp = u_print(idp, objA.a, objA.b)\n",
    "objA.f(10.)   # 100. * 10. = 1000.  expected, if \"a\" unchanged for \"f\" => 2.0  * 10. = 20."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "gvxqeOOixCdk"
   },
   "source": [
    "## We have to be careful:\n",
    "- the \"static\" character of \"a\" is a by-product of the analysis of the code by JIT, but this is done silently and without warning and therefore, this can mislead the user\n",
    "\n",
    "What is behind? The problem is that `static_argnums` relies on the hash method of the object to determine if it has changed between two calls, and the\n",
    "method `__hash__` for a user-defined class does not take into account the class attributes.\n",
    "\n",
    "This means that on the second function call, JAX has no way of knowing that the class attributes have changed and uses the static value cached during the previous compilation.\n",
    "\n",
    "For this reason, if you mark personal arguments as static, it is important that you define an appropriate `__hash__` method for your class. For example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "id": "CoCmcYuVw3hL"
   },
   "outputs": [],
   "source": [
    "class A():\n",
    "    def __init__(self, a):\n",
    "        print(\"new A\")\n",
    "        self.a = a\n",
    "        self.b = None\n",
    "\n",
    "    def set_a(self,x):\n",
    "        self.a = x\n",
    "        print(\"new a:\",self.a)\n",
    "\n",
    "\n",
    "    def set_b(self,x):\n",
    "        self.b = x\n",
    "\n",
    "    @partial(jit, static_argnums=0)   #\"self\" is static\n",
    "    def f(self, var):\n",
    "        print(\"compile...\")\n",
    "        return self.a * var\n",
    "\n",
    "    # specific hash metho\n",
    "    def __hash__(self):\n",
    "        return hash((self.a,self.b))\n",
    "\n",
    "    def __eq__(self, other):   # you need it also\n",
    "        return (isinstance(other, A) and\n",
    "            (self.a, self.b) == (other.a, other.b))\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qXA4m4EQz03c"
   },
   "source": [
    "But this is not the end...."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 869
    },
    "id": "DHM--OfjzkK5",
    "outputId": "5b00dbe0-ab28-4186-b8c6-8d8425440475"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "new A\n",
      "0): 2.0 None\n",
      "compile...\n",
      "b: 20.0 <class 'jaxlib.xla_extension.ArrayImpl'>\n",
      "1): 2.0 20.0\n"
     ]
    },
    {
     "ename": "ValueError",
     "evalue": "Non-hashable static arguments are not supported. An error occurred while trying to hash an object of type <class '__main__.A'>, <__main__.A object at 0x79424d9cba00>. The error was:\nTraceback (most recent call last):\n  File \"/usr/lib/python3.10/runpy.py\", line 196, in _run_module_as_main\n  File \"/usr/lib/python3.10/runpy.py\", line 86, in _run_code\n  File \"/usr/local/lib/python3.10/dist-packages/colab_kernel_launcher.py\", line 37, in <module>\n  File \"/usr/local/lib/python3.10/dist-packages/traitlets/config/application.py\", line 992, in launch_instance\n  File \"/usr/local/lib/python3.10/dist-packages/ipykernel/kernelapp.py\", line 619, in start\n  File \"/usr/local/lib/python3.10/dist-packages/tornado/platform/asyncio.py\", line 195, in start\n  File \"/usr/lib/python3.10/asyncio/base_events.py\", line 603, in run_forever\n  File \"/usr/lib/python3.10/asyncio/base_events.py\", line 1909, in _run_once\n  File \"/usr/lib/python3.10/asyncio/events.py\", line 80, in _run\n  File \"/usr/local/lib/python3.10/dist-packages/tornado/ioloop.py\", line 685, in <lambda>\n  File \"/usr/local/lib/python3.10/dist-packages/tornado/ioloop.py\", line 738, in _run_callback\n  File \"/usr/local/lib/python3.10/dist-packages/tornado/gen.py\", line 825, in inner\n  File \"/usr/local/lib/python3.10/dist-packages/tornado/gen.py\", line 786, in run\n  File \"/usr/local/lib/python3.10/dist-packages/ipykernel/kernelbase.py\", line 361, in process_one\n  File \"/usr/local/lib/python3.10/dist-packages/tornado/gen.py\", line 234, in wrapper\n  File \"/usr/local/lib/python3.10/dist-packages/ipykernel/kernelbase.py\", line 261, in dispatch_shell\n  File \"/usr/local/lib/python3.10/dist-packages/tornado/gen.py\", line 234, in wrapper\n  File \"/usr/local/lib/python3.10/dist-packages/ipykernel/kernelbase.py\", line 539, in execute_request\n  File \"/usr/local/lib/python3.10/dist-packages/tornado/gen.py\", line 234, in wrapper\n  File \"/usr/local/lib/python3.10/dist-packages/ipykernel/ipkernel.py\", line 302, in do_execute\n  File \"/usr/local/lib/python3.10/dist-packages/ipykernel/zmqshell.py\", line 539, in run_cell\n  File \"/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py\", line 2975, in run_cell\n  File \"/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py\", line 3030, in _run_cell\n  File \"/usr/local/lib/python3.10/dist-packages/IPython/core/async_helpers.py\", line 78, in _pseudo_sync_runner\n  File \"/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py\", line 3257, in run_cell_async\n  File \"/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py\", line 3473, in run_ast_nodes\n  File \"/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py\", line 3553, in run_code\n  File \"<ipython-input-29-ac13534385f8>\", line 11, in <cell line: 11>\n  File \"<ipython-input-28-7e9f3e1bd40e>\", line 22, in __hash__\nTypeError: unhashable type: 'jaxlib.xla_extension.DeviceList'\n",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-29-ac13534385f8>\u001b[0m in \u001b[0;36m<cell line: 11>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      9\u001b[0m \u001b[0midp\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mu_print\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0midp\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mobjA\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0ma\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mobjA\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mb\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     10\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 11\u001b[0;31m \u001b[0mobjA\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mf\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m11.\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m: Non-hashable static arguments are not supported. An error occurred while trying to hash an object of type <class '__main__.A'>, <__main__.A object at 0x79424d9cba00>. The error was:\nTraceback (most recent call last):\n  File \"/usr/lib/python3.10/runpy.py\", line 196, in _run_module_as_main\n  File \"/usr/lib/python3.10/runpy.py\", line 86, in _run_code\n  File \"/usr/local/lib/python3.10/dist-packages/colab_kernel_launcher.py\", line 37, in <module>\n  File \"/usr/local/lib/python3.10/dist-packages/traitlets/config/application.py\", line 992, in launch_instance\n  File \"/usr/local/lib/python3.10/dist-packages/ipykernel/kernelapp.py\", line 619, in start\n  File \"/usr/local/lib/python3.10/dist-packages/tornado/platform/asyncio.py\", line 195, in start\n  File \"/usr/lib/python3.10/asyncio/base_events.py\", line 603, in run_forever\n  File \"/usr/lib/python3.10/asyncio/base_events.py\", line 1909, in _run_once\n  File \"/usr/lib/python3.10/asyncio/events.py\", line 80, in _run\n  File \"/usr/local/lib/python3.10/dist-packages/tornado/ioloop.py\", line 685, in <lambda>\n  File \"/usr/local/lib/python3.10/dist-packages/tornado/ioloop.py\", line 738, in _run_callback\n  File \"/usr/local/lib/python3.10/dist-packages/tornado/gen.py\", line 825, in inner\n  File \"/usr/local/lib/python3.10/dist-packages/tornado/gen.py\", line 786, in run\n  File \"/usr/local/lib/python3.10/dist-packages/ipykernel/kernelbase.py\", line 361, in process_one\n  File \"/usr/local/lib/python3.10/dist-packages/tornado/gen.py\", line 234, in wrapper\n  File \"/usr/local/lib/python3.10/dist-packages/ipykernel/kernelbase.py\", line 261, in dispatch_shell\n  File \"/usr/local/lib/python3.10/dist-packages/tornado/gen.py\", line 234, in wrapper\n  File \"/usr/local/lib/python3.10/dist-packages/ipykernel/kernelbase.py\", line 539, in execute_request\n  File \"/usr/local/lib/python3.10/dist-packages/tornado/gen.py\", line 234, in wrapper\n  File \"/usr/local/lib/python3.10/dist-packages/ipykernel/ipkernel.py\", line 302, in do_execute\n  File \"/usr/local/lib/python3.10/dist-packages/ipykernel/zmqshell.py\", line 539, in run_cell\n  File \"/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py\", line 2975, in run_cell\n  File \"/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py\", line 3030, in _run_cell\n  File \"/usr/local/lib/python3.10/dist-packages/IPython/core/async_helpers.py\", line 78, in _pseudo_sync_runner\n  File \"/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py\", line 3257, in run_cell_async\n  File \"/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py\", line 3473, in run_ast_nodes\n  File \"/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py\", line 3553, in run_code\n  File \"<ipython-input-29-ac13534385f8>\", line 11, in <cell line: 11>\n  File \"<ipython-input-28-7e9f3e1bd40e>\", line 22, in __hash__\nTypeError: unhashable type: 'jaxlib.xla_extension.DeviceList'\n"
     ]
    }
   ],
   "source": [
    "idp=0\n",
    "\n",
    "objA = A(2.0)\n",
    "idp = u_print(idp, objA.a, objA.b)\n",
    "\n",
    "b = objA.f(10.)         #  what is the type of 'b' ??? (*)\n",
    "print(\"b:\", b, type(b))\n",
    "objA.set_b(b)\n",
    "idp = u_print(idp, objA.a, objA.b)\n",
    "\n",
    "objA.f(11.)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "BAeHhF6K0FDe"
   },
   "source": [
    "the problem comes from the fact that after compilation `b` becomes a `DeviceArray` which is unhastable so it's stuck because then self can't be hashtable too !!!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "id": "ekWiOl--z41k"
   },
   "outputs": [],
   "source": [
    "class A():\n",
    "    def __init__(self, a):\n",
    "        print(\"Nouveau A\")\n",
    "        self.a = a\n",
    "        self.b = None\n",
    "\n",
    "    def set_a(self,x):\n",
    "        self.a = x\n",
    "        print(\"new a:\",self.a)\n",
    "\n",
    "    def set_b(self,x):\n",
    "        self.b = x\n",
    "\n",
    "    @partial(jit, static_argnums=0)\n",
    "    def f(self, var):\n",
    "        print(\"compile...\")\n",
    "        return self.a * var\n",
    "\n",
    "    def __hash__(self):\n",
    "        return hash((self.a))                 # do not put \"b\"\n",
    "\n",
    "    def __eq__(self, other):\n",
    "        return (isinstance(other, A) and\n",
    "            (self.a) == (other.a))             # do not put \"b\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "pqckfEKQ0ZJY",
    "outputId": "4caf2ddc-3d87-4815-82b7-56779dbb2839"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Nouveau A\n",
      "0): 2.0 None\n",
      "compile...\n",
      "b: 20.0 <class 'jaxlib.xla_extension.ArrayImpl'>\n",
      "1): 2.0 20.0\n",
      "2): 2.0 22.0\n",
      "Nouveau A\n",
      "3): 3.0 None\n",
      "compile...\n",
      "4): 3.0 60.0\n"
     ]
    }
   ],
   "source": [
    "idp=0\n",
    "\n",
    "objA = A(2.0)\n",
    "idp = u_print(idp, objA.a, objA.b)\n",
    "\n",
    "b = objA.f(10.)         #  quel est le type de 'b' ??? (*)\n",
    "print(\"b:\", b, type(b))\n",
    "objA.set_b(b)\n",
    "idp = u_print(idp, objA.a, objA.b)\n",
    "\n",
    "objA.set_b(objA.f(11.))\n",
    "idp = u_print(idp, objA.a, objA.b)\n",
    "\n",
    "\n",
    "objA = A(3.0)\n",
    "idp = u_print(idp, objA.a, objA.b)\n",
    "objA.set_b(objA.f(20.))\n",
    "idp = u_print(idp, objA.a, objA.b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "nHTbHn3b0nDL",
    "outputId": "05f61445-22e4-4615-c043-98cb247a9fbf"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "new a: 40.0\n",
      "compile...\n",
      "5): 40.0 440.0\n"
     ]
    }
   ],
   "source": [
    "objA.set_a(40.)\n",
    "objA.set_b(objA.f(11.))\n",
    "idp = u_print(idp, objA.a, objA.b)   # 40.*11. = 440\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "LDNbI22W001r"
   },
   "source": [
    "## Review self as static:\n",
    "\n",
    "- We have to carefully define `__hass__`  and `__eq__` (see Python [doc](https://docs.python.org/3/reference/datamodel.html#object.__hash__)) but DeviceArray are non-hashtable\n",
    "\n",
    "- `a` can be redefined voluntarily (with/without setter), the result of `b` is well updated\n",
    "\n",
    "- So can be done but with care and there is a simpler solution\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Dj-Nxvuk1ync"
   },
   "source": [
    "# 3d strategy: PyTree"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "id": "IjIsNeyM0b94"
   },
   "outputs": [],
   "source": [
    "class A():\n",
    "    def __init__(self, a, b=None):   # new signature\n",
    "        #print(\"Nouveau A\")\n",
    "        self.a = a\n",
    "        self.b = b\n",
    "\n",
    "\n",
    "    def set_a(self,x):\n",
    "        self.a = x\n",
    "        print(\"new a:\",self.a)\n",
    "\n",
    "    def set_b(self, b):\n",
    "        self.b = b\n",
    "\n",
    "    @jit                              # <------ self, no more static    !\n",
    "    def f(self, var):\n",
    "        print(\"compile...\")\n",
    "        return self.a * var\n",
    "\n",
    "    def g(self):\n",
    "      return self.b * self.b\n",
    "\n",
    "    #### PyTree methods....\n",
    "    def tree_flatten(self):\n",
    "        children = (self.b,)         # dynamic values\n",
    "        aux_data = {'a': self.a}     # static values\n",
    "        return (children, aux_data)\n",
    "\n",
    "    @classmethod\n",
    "    def tree_unflatten(cls, aux_data, children):\n",
    "        b = children\n",
    "        a = aux_data['a']\n",
    "        return cls(a=a, b=b ) # should be the signature of __init__\n",
    "\n",
    "# register explicit Or decorator (as Cosmology class)\n",
    "from jax import tree_util\n",
    "tree_util.register_pytree_node(A,\n",
    "                               A.tree_flatten,\n",
    "                               A.tree_unflatten)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "ILNoH3042Gkm",
    "outputId": "9df8b6c9-b59a-464c-b38a-a41121b6d460"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0): 2.0 None\n",
      "compile...\n",
      "b: 20.0 <class 'jaxlib.xla_extension.ArrayImpl'>\n",
      "1): 2.0 20.0\n",
      "compile...\n",
      "2): 2.0 22.0\n",
      "3): 3.0 None\n",
      "compile...\n",
      "4): 3.0 60.0\n",
      "compile...\n",
      "5): 3.0 90.0\n",
      "new a: 30.0\n",
      "compile...\n",
      "6): 30.0 900.0\n",
      "7): 3.0 None\n",
      "8): 3.0 60.0\n",
      "9): 3.0 90.0\n",
      ">>>> try an array for var\n",
      "compile...\n",
      "10): 3.0 [90.]\n",
      "compile...\n",
      "11): 3.0 [120.]\n",
      "g:  [14400.]\n"
     ]
    }
   ],
   "source": [
    "idp=0\n",
    "\n",
    "objA = A(2.0)\n",
    "idp = u_print(idp, objA.a, objA.b)\n",
    "\n",
    "b = objA.f(10.)         #  quel est le type de 'b' ??? (*)\n",
    "print(\"b:\", b, type(b))\n",
    "objA.set_b(b)\n",
    "idp = u_print(idp, objA.a, objA.b)\n",
    "\n",
    "objA.set_b(objA.f(11.))\n",
    "idp = u_print(idp, objA.a, objA.b)\n",
    "\n",
    "\n",
    "objA = A(3.0)\n",
    "idp = u_print(idp, objA.a, objA.b)\n",
    "objA.set_b(objA.f(20.))\n",
    "idp = u_print(idp, objA.a, objA.b)\n",
    "objA.set_b(objA.f(30.))\n",
    "idp = u_print(idp, objA.a, objA.b)\n",
    "\n",
    "\n",
    "objA.set_a(30.)\n",
    "objA.set_b(objA.f(30.))\n",
    "idp = u_print(idp, objA.a, objA.b)\n",
    "\n",
    "\n",
    "\n",
    "objA = A(3.0)\n",
    "idp = u_print(idp, objA.a, objA.b)\n",
    "\n",
    "objA.set_b(objA.f(20.))\n",
    "idp = u_print(idp, objA.a, objA.b)\n",
    "\n",
    "objA.set_b(objA.f(30.))\n",
    "idp = u_print(idp, objA.a, objA.b)\n",
    "\n",
    "print(\">>>> try an array for var\")\n",
    "\n",
    "objA.set_b(objA.f(jnp.array([30.])))\n",
    "idp = u_print(idp, objA.a, objA.b)\n",
    "\n",
    "objA.set_b(objA.f(jnp.array([40.])))\n",
    "idp = u_print(idp, objA.a, objA.b)\n",
    "\n",
    "print(\"g: \", objA.g())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wWbd0Zeg2lwS"
   },
   "source": [
    "## A possible variation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {
    "id": "8J4MwGRl2KqP"
   },
   "outputs": [],
   "source": [
    "class A():\n",
    "    def __init__(self, a, b=None):\n",
    "        self.a = a\n",
    "        self.b = b\n",
    "\n",
    "    @jit\n",
    "    def f(self, var):\n",
    "        print(\"compile...\")\n",
    "        new_b = self.a * var\n",
    "        return A(self.a, new_b)       # <--------- new objet.\n",
    "\n",
    "\n",
    "# registration with  lambda functions\n",
    "tree_util.register_pytree_node(A,\n",
    "                               lambda x: ((x.a,x.b), None),\n",
    "                               lambda _, x: A(a=x[0],b=x[1])\n",
    "                              )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "qRW7lVXv2uSb",
    "outputId": "91403688-5475-439d-b76f-f1043dd4e39e"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0): 2.0 None\n",
      "compile...\n",
      "1): 2.0 20.0\n",
      "compile...\n",
      "2): 2.0 22.0\n",
      "3): 3.0 None\n",
      "4): 3.0 60.0\n",
      "5): 3.0 90.0\n",
      "6): 30.0 900.0\n",
      "7): 3.0 None\n",
      "8): 3.0 60.0\n",
      "9): 3.0 90.0\n",
      ">>>> try an array for var\n",
      "compile...\n",
      "10): 3.0 [90.]\n",
      "compile...\n",
      "11): 3.0 [120.]\n",
      "12): 3.0 [150.]\n",
      "13): 40.0 [2000.]\n"
     ]
    }
   ],
   "source": [
    "idp=0\n",
    "\n",
    "objA = A(2.0)\n",
    "idp = u_print(idp, objA.a, objA.b)    #0\n",
    "\n",
    "objA = objA.f(10.)\n",
    "idp = u_print(idp,objA.a, objA.b)     #1\n",
    "\n",
    "objA = objA.f(11.)\n",
    "idp = u_print(idp, objA.a, objA.b)    #2\n",
    "\n",
    "####\n",
    "\n",
    "objA = A(3.0)\n",
    "idp = u_print(idp, objA.a, objA.b)     #3\n",
    "objA = objA.f(20.)\n",
    "idp = u_print(idp, objA.a, objA.b)     #4\n",
    "objA= objA.f(30.)\n",
    "idp = u_print(idp, objA.a, objA.b)     #5\n",
    "\n",
    "\n",
    "objA.a = 30.  # change\n",
    "objA= objA.f(30.)\n",
    "idp = u_print(idp, objA.a, objA.b)     #6\n",
    "\n",
    "####\n",
    "\n",
    "objA = A(3.0)\n",
    "idp = u_print(idp, objA.a, objA.b)     #7\n",
    "\n",
    "objA= objA.f(20.)\n",
    "idp = u_print(idp, objA.a, objA.b)     #8\n",
    "\n",
    "objA= objA.f(30.)\n",
    "idp = u_print(idp, objA.a, objA.b)     #9\n",
    "\n",
    "print(\">>>> try an array for var\")\n",
    "\n",
    "objA= objA.f(jnp.array([30.]))\n",
    "idp = u_print(idp, objA.a, objA.b)     #10\n",
    "\n",
    "objA= objA.f(jnp.array([40.]))\n",
    "idp = u_print(idp, objA.a, objA.b)     #11\n",
    "\n",
    "objA= objA.f(jnp.array([50.]))\n",
    "idp = u_print(idp, objA.a, objA.b)     #12\n",
    "\n",
    "\n",
    "objA.a = 40. ### new change\n",
    "objA= objA.f(jnp.array([50.]))\n",
    "idp = u_print(idp, objA.a, objA.b)     #13\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "AtnQOf9S3GAs"
   },
   "source": [
    "## Review PyTree\n",
    "\n",
    "- one must somehow set up `tree_flatten` and `tree_unflatten` with care\n",
    "- the change of ``a` is well transmitted\n",
    "- choice of implementation with new object each time or setter of the internal variable to change"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {
    "id": "V-DbnwUV2zhz"
   },
   "outputs": [],
   "source": [
    "from jax.tree_util import register_pytree_node_class\n",
    "\n",
    "@register_pytree_node_class\n",
    "class B():\n",
    "    def __init__(self, p, v):\n",
    "        self.p = p\n",
    "        self.v = v\n",
    "        print(\"init:\",p,v)\n",
    "\n",
    "    @jit\n",
    "    def step(self, dt):\n",
    "        print(\"compile...\")\n",
    "        a = -9.8\n",
    "        new_v = self.v + a * dt\n",
    "        new_p = self.p + new_v * dt\n",
    "        return B(new_p, new_v)\n",
    "\n",
    "    #### PyTree methods....\n",
    "    def tree_flatten(self):\n",
    "        print(\"tree_flatten\")\n",
    "        children = (self.p,self.v)         # dynamic values\n",
    "        aux_data = {}     # static values\n",
    "        return (children, aux_data)\n",
    "\n",
    "    @classmethod\n",
    "    def tree_unflatten(cls, aux_data, children):\n",
    "        print(\"tree_unflatten\")\n",
    "        b = children\n",
    "        return cls(p=b[0], v=b[1] ) # should be the signature of __init__\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "K9NHO7Op4j9B",
    "outputId": "4f5273ad-3cd4-4b1f-8a3a-3433dc08a9d5"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "init: [0.  0.5] [1.  1.5]\n"
     ]
    }
   ],
   "source": [
    "b = B(jnp.array([0., 0.5]), jnp.array([1., 1.5]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "l-eTQ0KQ5ZPn",
    "outputId": "f2796cd5-ed09-4e41-da01-8a568de5b6b6"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tree_flatten\n",
      "tree_flatten\n",
      "tree_flatten\n",
      "tree_unflatten\n",
      "init: Traced<ShapedArray(float64[2])>with<DynamicJaxprTrace(level=1/0)> Traced<ShapedArray(float64[2])>with<DynamicJaxprTrace(level=1/0)>\n",
      "compile...\n",
      "init: Traced<ShapedArray(float64[2])>with<DynamicJaxprTrace(level=1/0)> Traced<ShapedArray(float64[2])>with<DynamicJaxprTrace(level=1/0)>\n",
      "tree_flatten\n",
      "tree_flatten\n",
      "tree_unflatten\n",
      "init: [0.00902 0.51402] [0.902 1.402]\n",
      "tree_unflatten\n",
      "init: [0.00902 0.51402] [0.902 1.402]\n",
      "tree_flatten\n",
      "tree_flatten\n",
      "tree_unflatten\n",
      "init: [0.01706 0.52706] [0.804 1.304]\n",
      "tree_flatten\n",
      "tree_unflatten\n",
      "init: [0.02412 0.53912] [0.706 1.206]\n",
      "tree_flatten\n",
      "tree_unflatten\n",
      "init: [0.0302 0.5502] [0.608 1.108]\n",
      "tree_flatten\n",
      "tree_unflatten\n",
      "init: [0.0353 0.5603] [0.51 1.01]\n"
     ]
    }
   ],
   "source": [
    "for i in range(5):\n",
    "    b = b.step(0.01)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "iHRhSOZL4dWv",
    "outputId": "fb282140-afbf-42a7-e5b7-bd27453ddbc3"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m175.2/175.2 kB\u001b[0m \u001b[31m1.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m41.2/41.2 kB\u001b[0m \u001b[31m4.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25h"
     ]
    }
   ],
   "source": [
    "! pip install -q equinox"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {
    "id": "RmdWjOGU_kdh"
   },
   "outputs": [],
   "source": [
    "import equinox as eqx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {
    "id": "ehcPIqsK_fM4"
   },
   "outputs": [],
   "source": [
    "class D(eqx.Module):\n",
    "  p: float\n",
    "  v: float\n",
    "\n",
    "  def __init__(self, p,v):  # the init\n",
    "    self.p = p\n",
    "    self.v = v\n",
    "\n",
    "  def __call__(self, dt):    # the forward call\n",
    "        a = -9.8\n",
    "        new_v = self.v + a * dt\n",
    "        new_p = self.p + new_v * dt\n",
    "        return D(new_p, new_v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {
    "id": "k6qop4n5BA64"
   },
   "outputs": [],
   "source": [
    "d = D(np.array([0, 0.5]), np.array([1, 1.5]))\n",
    "for i in range(5):\n",
    "    d = d(0.01)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "cOW0pcmtBxWu",
    "outputId": "14d4d2d2-8fb9-47fc-bbee-0118891dc9d9"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([0.0353, 0.5603]), array([0.51, 1.01]))"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "d.p, d.v"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "izCf_wEACEj1"
   },
   "source": [
    "# Takeaway message (JIT in a class):\n",
    "\n",
    "- several methods: helper, self-static, Pytree\n",
    "\n",
    "- should not use `self.<var>` = in a jitted function (otherwise side effect). By the way, the code analysis by JIT is done via `Traced<ShapedArray>`\n",
    "\n",
    "- use of `__hash__` and `__eq__`, and the DeviceArray is unhastable\n",
    "\n",
    "- when the compilation is triggered (JIT)\n",
    "\n",
    "- See also the doc [to-jit-or-not-to-jit](https://jax.readthedocs.io/en/latest/notebooks/thinking_in_jax.html?highlight=Traced%3CShapedArray%3E#to-jit-or-not-to-jit)\n",
    "\n",
    "- Ask yourself \"Do I need OO encapsulation (i.e. make Class) ?\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "3iN6Bg6kAhXR"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "authorship_tag": "ABX9TyPe4VPLXXoIQitUdwJg4vPb",
   "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
}