JAX_PyTree_initialisation.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_PyTree_initialisation.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": "awI0tvPtkd-k"
   },
   "outputs": [],
   "source": [
    "import jax\n",
    "from jax import lax, grad, jit, vmap\n",
    "from jax.tree_util import register_pytree_node_class\n",
    "\n",
    "import jax.numpy as jnp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 36
    },
    "id": "WVBuuMNbDfn6",
    "outputId": "174f4d98-a7c9-40ef-f48f-63e95355f635"
   },
   "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": "markdown",
   "metadata": {
    "id": "z8LJTNRsknZc"
   },
   "source": [
    "# Topic: more on PyTrees\n",
    "\n",
    " - When tree_flatten/tree_unflatten are triggered?\n",
    " - Initialisation that crashes: Why? a solution\n",
    " - More advenced with different kinds of variable usage"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "NpTL3uiYrMaM"
   },
   "source": [
    "## When tree_flatten/tree_unflatten are triggered <=> statefull initialisation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "id": "1PYK22yvklZl"
   },
   "outputs": [],
   "source": [
    "jax.clear_caches()\n",
    "\n",
    "@register_pytree_node_class\n",
    "class State:\n",
    "  def __init__(self, n: int):\n",
    "    self.n = n\n",
    "    self.stateful()\n",
    "    jax.debug.print('   n = {} # init', self.n)  # this print the value not the kind of Array Traced by JAX\n",
    "\n",
    "  def stateful(self):\n",
    "    self.n += 100\n",
    "\n",
    "  def tree_flatten(self):\n",
    "    jax.debug.print('   n = {} # flatten', self.n + 10)\n",
    "    return (self.n,), {}\n",
    "\n",
    "  @classmethod\n",
    "  def tree_unflatten(cls, aux_data, children):\n",
    "    n, = children\n",
    "    jax.debug.print('   n = {} # unflatten', n + 1)\n",
    "    return cls(n)  # runs `__init__` and so call `stateful`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "p8CzvEKjngTA"
   },
   "source": [
    "As a reminder jax.lax.scan is equivalent of\n",
    "```python\n",
    "def scan(f, init, xs, length=None):\n",
    "  if xs is None:\n",
    "    xs = [None] * length\n",
    "  carry = init\n",
    "  ys = []\n",
    "  for x in xs:\n",
    "    carry, y = f(carry, x)\n",
    "    ys.append(y)\n",
    "  return carry, np.stack(ys)\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "3SBaQouElBXj",
    "outputId": "998ca22d-a9c8-4b4e-cccb-3a4c24dfe8a4"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "   n = 100 # init\n",
      "%%%\n",
      "xs: [0 1 2]\n",
      "xxx\n",
      "   n = 110 # flatten\n",
      "   n = 110 # flatten\n",
      "   n = 110 # flatten\n",
      "   n = 110 # flatten\n",
      "   n = 101 # unflatten\n",
      "   n = 200 # init\n",
      "---\n",
      "0: n = 200 # body\n",
      "   n = 210 # flatten\n",
      "   n = 201 # unflatten\n",
      "   n = 300 # init\n",
      "---\n",
      "1: n = 300 # body\n",
      "   n = 310 # flatten\n",
      "   n = 301 # unflatten\n",
      "   n = 400 # init\n",
      "---\n",
      "2: n = 400 # body\n",
      "   n = 410 # flatten\n",
      "   n = 401 # unflatten\n",
      "   n = 500 # init\n",
      "y:  [200 300 400]\n"
     ]
    }
   ],
   "source": [
    "def body_fun(carry, x):\n",
    "  jax.debug.print('---')\n",
    "  jax.debug.print('{}: n = {} # body', x, carry.n)\n",
    "  return carry, carry.n   # carry.n is the \"y\" in the python-equivalent code\n",
    "\n",
    "s = State(0)  # we expect n=100\n",
    "print(\"%%%\")\n",
    "xs = jnp.arange(3, dtype=jnp.int16)\n",
    "print(\"xs:\", xs)\n",
    "print('xxx')\n",
    "# s is the \"carry\"\n",
    "_, y = lax.scan(body_fun, s, xs)   # scan will first compile body_fun\n",
    "print(\"y: \",y)  # what do you expect ? [100, 100, 100], [100, 200, 300]  or something else"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vJ6NNpfYnHtk"
   },
   "source": [
    "You see that : this is the body output that is stored , but the initialisation is called before and so the state of the object is silently changed before the scan is performed."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "H652hpxDrU0v"
   },
   "source": [
    "## Initialisation of a PyTree"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZtYO3UKHrdTt"
   },
   "source": [
    "### Let us start with a crash"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "id": "_KlwKeEqlQ77"
   },
   "outputs": [],
   "source": [
    "class BaseObj():\n",
    "    def __init__(self, *, gsparams=None, **params):\n",
    "      self._params = params      # Dictionary containing all traced parameters\n",
    "      self._gsparams = gsparams  # Non-traced static parameters\n",
    "      jax.debug.print('Init BaseObj')\n",
    "\n",
    "\n",
    "    @property\n",
    "    def gsparams(self):\n",
    "        \"\"\"A `GSParams` object that sets various parameters relevant for speed/accuracy trade-offs.\"\"\"\n",
    "        return self._gsparams\n",
    "\n",
    "    @property\n",
    "    def params(self):\n",
    "        \"\"\"A Dictionary object containing all parameters of the internal represention of this object.\"\"\"\n",
    "        return self._params\n",
    "\n",
    "\n",
    "    def tree_flatten(self):\n",
    "        \"\"\"This function flattens the BaseObj into a list of children\n",
    "        nodes that will be traced by JAX and auxiliary static data.\"\"\"\n",
    "        jax.debug.print('tree_flatten')\n",
    "        # Define the children nodes of the PyTree that need tracing\n",
    "        children = (self.params,)\n",
    "        # Define auxiliary static data that doesn’t need to be traced\n",
    "        aux_data = {\"gsparams\": self.gsparams}\n",
    "        return (children, aux_data)\n",
    "\n",
    "    @classmethod\n",
    "    def tree_unflatten(cls, aux_data, children):\n",
    "        \"\"\"Recreates an instance of the class from flatten representation\"\"\"\n",
    "        jax.debug.print('tree_unflatten')\n",
    "        return cls(**(children[0]), **aux_data)\n",
    "\n",
    "@register_pytree_node_class\n",
    "class T(BaseObj):\n",
    "  def __init__(self, a, gsparams=None):\n",
    "    super().__init__(a=a, gsparams=gsparams)\n",
    "\n",
    "    jax.debug.print(\"Init T: a={}, type {}\",a, type(a))\n",
    "    print(\"a = \",a)\n",
    "    self._m = a*2     #<<<<<<<<<<<<<<  pb if decommented when we use a vmap\n",
    "\n",
    "\n",
    "  @property\n",
    "  def a(self):\n",
    "    return self.params[\"a\"]\n",
    "\n",
    "  def f(self,x):\n",
    "    return self.a*x + self._m"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PUZD0lq7r8ow"
   },
   "source": [
    "If I instantiate a collection of T objs and apply a vmap"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "3aFrEmrVr3yv",
    "outputId": "ed381f7b-da35-4475-857c-e3bf8303841a"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Init BaseObj\n",
      "Init T: a=[3.1415927 6.2831855], type <class 'jaxlib.xla_extension.ArrayImpl'>\n",
      "a =  [3.1415927 6.2831855]\n"
     ]
    }
   ],
   "source": [
    "avals = jnp.pi * jnp.array([1.,2.])\n",
    "ts = T(avals)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 560
    },
    "id": "kb_uJy1er_nx",
    "outputId": "8fa11cc4-8f51-4799-c2a1-9d0b5fcd99ae"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tree_flatten\n",
      "tree_unflatten\n",
      "Init BaseObj\n",
      "Init T: a=<object object at 0x7ff4fcfdef40>, type <class 'object'>\n",
      "a =  <object object at 0x7ff4fcfdef40>\n"
     ]
    },
    {
     "ename": "TypeError",
     "evalue": "unsupported operand type(s) for *: 'object' and 'int'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-17-06b641f62fff>\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[0mvmap\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;32mlambda\u001b[0m \u001b[0mt\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mt\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mf\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[0mts\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 4 frame]\u001b[0m\n",
      "\u001b[0;32m<ipython-input-15-19164efc7592>\u001b[0m in \u001b[0;36mtree_unflatten\u001b[0;34m(cls, aux_data, children)\u001b[0m\n\u001b[1;32m     31\u001b[0m         \u001b[0;34m\"\"\"Recreates an instance of the class from flatten representation\"\"\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     32\u001b[0m         \u001b[0mjax\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdebug\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'tree_unflatten'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 33\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mcls\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m**\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mchildren\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[0maux_data\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     34\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     35\u001b[0m \u001b[0;34m@\u001b[0m\u001b[0mregister_pytree_node_class\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-15-19164efc7592>\u001b[0m in \u001b[0;36m__init__\u001b[0;34m(self, a, gsparams)\u001b[0m\n\u001b[1;32m     40\u001b[0m     \u001b[0mjax\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdebug\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"Init T: a={}, type {}\"\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0ma\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtype\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0ma\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     41\u001b[0m     \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"a = \"\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0ma\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 42\u001b[0;31m     \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_m\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0ma\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0;36m2\u001b[0m     \u001b[0;31m#<<<<<<<<<<<<<<  pb if decommented when we use a vmap\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     43\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     44\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mTypeError\u001b[0m: unsupported operand type(s) for *: 'object' and 'int'"
     ]
    }
   ],
   "source": [
    "jax.vmap(lambda t: t.f(1.))(ts)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9Asf6kWqt6hO"
   },
   "source": [
    "As you see `a` is Traced object, so it does not make sense to perform any operation with it. Notice that the crash doesn't depends of the \"super\" init position.\n",
    "\n",
    "### One simple solution in that case"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "id": "yK7FzMOntQQM"
   },
   "outputs": [],
   "source": [
    "@register_pytree_node_class\n",
    "class T(BaseObj):\n",
    "  def __init__(self, a, gsparams=None):\n",
    "    super().__init__(a=a, gsparams=gsparams)\n",
    "\n",
    "    jax.debug.print(\"Init T: a={}\",a)\n",
    "\n",
    "\n",
    "  @property\n",
    "  def _m(self):                #<-----\n",
    "    return self.a * 2\n",
    "\n",
    "  @property\n",
    "  def a(self):\n",
    "    return self.params[\"a\"]\n",
    "\n",
    "  @jit\n",
    "  def f(self,x):\n",
    "    return self.a*x + self._m"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "56lXz9ZMv0wx",
    "outputId": "f2336d29-26c8-47f6-c6a4-2abb4946dfd2"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Init BaseObj\n",
      "Init T: a=[3.1415927 6.2831855]\n"
     ]
    }
   ],
   "source": [
    "avals = jnp.pi * jnp.array([1.,2.])\n",
    "ts = T(avals)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "26jMLXrhv3-v",
    "outputId": "ff1a2978-7129-4b01-9d5d-c766fec74a8b"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tree_flatten\n",
      "tree_unflatten\n",
      "Init BaseObj\n",
      "Init T: a=<object object at 0x7ff4fcfdf590>\n",
      "tree_flatten\n",
      "tree_unflatten\n",
      "Init BaseObj\n",
      "Init T: a=3.1415927410125732\n",
      "Init T: a=6.2831854820251465\n",
      "tree_flatten\n",
      "tree_flatten\n",
      "tree_flatten\n",
      "tree_unflatten\n",
      "Init BaseObj\n",
      "Init T: a=3.1415927410125732\n",
      "Init T: a=6.2831854820251465\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Array([ 9.424778, 18.849556], dtype=float32)"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jax.vmap(lambda t: t.f(1.))(ts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "rQXr-1jjx092",
    "outputId": "6c5d4bb0-241f-4522-ca92-a6f1b8aed9ea"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tree_flatten\n",
      "tree_unflatten\n",
      "Init BaseObj\n",
      "Init T: a=<object object at 0x7ff4fcfde730>\n",
      "tree_flatten\n",
      "tree_unflatten\n",
      "Init BaseObj\n",
      "Init T: a=3.1415927410125732\n",
      "Init T: a=6.2831854820251465\n",
      "tree_flatten\n",
      "tree_flatten\n",
      "tree_flatten\n",
      "tree_unflatten\n",
      "Init BaseObj\n",
      "Init T: a=3.1415927410125732\n",
      "Init T: a=6.2831854820251465\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Array([12.566371, 25.132742], dtype=float32)"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jax.vmap(lambda t: t.f(2.))(ts)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xB_NYjfww8zk"
   },
   "source": [
    "It works as expected but may be not convenient if the computation of `_m` in the above example needs a more sophisticated plan. Ie. we do not want to repeate it if for instance an other function need it."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Oyb_VchpzO-_"
   },
   "source": [
    "## A more complete ex with different kinds of variables\n",
    "\n",
    "The classes below are just for academic purposes to show the init & use of some variables that you can encontered in your use-case.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "id": "_7Rdl9DZv8gs"
   },
   "outputs": [],
   "source": [
    "from functools import partial\n",
    "\n",
    "from typing import Optional\n",
    "import jax.tree_util as tu\n",
    "from jax import Array, jit, lax, vmap\n",
    "from jax.typing import ArrayLike\n",
    "\n",
    "import numpy as np\n",
    "import numpy.typing as npt\n",
    "rng = np.random.default_rng(0)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "F2O7xAtJzIkx"
   },
   "source": [
    "### A Numpy class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "id": "K8u0C6qAyyoa"
   },
   "outputs": [],
   "source": [
    "class ExprNp:\n",
    "  def __init__(self, name: str, steps: int, const: npt.NDArray,\n",
    "               *, start: int = 0):\n",
    "    self.name = name  # metadata\n",
    "    self.steps = steps  # constant, static\n",
    "    self.const = const  # constant, array\n",
    "    # state, initialized in __init__, modified throughout\n",
    "    self._start = start\n",
    "    self._state = np.zeros_like(const, dtype=np.float32)\n",
    "    # value, depends on constant, initialized in __init__, never modified again\n",
    "    self._y = self._init_y(const)\n",
    "    # state, depends on state, initialized in __init__, modified throughout\n",
    "    self._sum = np.zeros(steps)\n",
    "\n",
    "  @property\n",
    "  def state(self):\n",
    "    return self._state , self._sum\n",
    "\n",
    "  @staticmethod\n",
    "  def _init_y(c: npt.NDArray):\n",
    "    return np.sin(c)\n",
    "\n",
    "  def _step(self, i):\n",
    "    self._state += self._y\n",
    "    self._sum[i] += 1\n",
    "\n",
    "  def main(self, steps: Optional[int] = None, *, start: Optional[int] = None):\n",
    "    steps = steps or self.steps\n",
    "    start = start or self._start\n",
    "    end = start + steps\n",
    "    for i in range(start, end):\n",
    "      self._step(i)\n",
    "    self._start = end\n",
    "    return self"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "68oW_yxtRjo4"
   },
   "source": [
    "### The JAX class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "2KxMbDNCSJG0",
    "outputId": "614c9126-e586-42aa-9f08-a089fb9fdadc"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.]\n"
     ]
    }
   ],
   "source": [
    "#small remark\n",
    "a = None or jnp.array([0.])\n",
    "print(a)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "id": "qXYIkD7g0F53"
   },
   "outputs": [],
   "source": [
    "@register_pytree_node_class\n",
    "class ExprJax:\n",
    "  def __init__(self, name: str, steps: int, const: Array, start: Array,\n",
    "               _state: Array, _sum: Array, _y: Array):\n",
    "    #\n",
    "    # the user should use the explicit init\n",
    "    #\n",
    "\n",
    "    # init in __init__, not changed after\n",
    "    self.name = name    # metadata here a string\n",
    "    self.steps = steps  # constant, static\n",
    "    self.const = const  # constant, array\n",
    "    # state, initialized in __init__, modified throughout\n",
    "    self._start = start\n",
    "    self._state = _state\n",
    "    # state, depends on state, initialized in __init__, modified throughout\n",
    "    self._sum = _sum\n",
    "    self._y = _y\n",
    "\n",
    "  # user initialisation\n",
    "  @classmethod\n",
    "  def init(cls, name: str, steps: int, const: ArrayLike, *, start: int = 0):\n",
    "    _state = jnp.zeros_like(const, dtype=jnp.float32)\n",
    "    _start = jnp.full(_state.shape[:-1], start)\n",
    "    _sum = jnp.zeros(_state.shape[:-1] + (steps,))\n",
    "    _y = cls._y_init(const)  # <<< initilize `y` in the smart constructor\n",
    "    return cls(name, steps, jnp.array(const), _start, _state=_state, _sum=_sum, _y=_y)\n",
    "\n",
    "  @property\n",
    "  def state(self):\n",
    "    return self._state , self._sum\n",
    "\n",
    "  def tree_flatten(self):\n",
    "    # <<< also passing `y` to `tree_flatten` to make sure it's not re-initialized\n",
    "    traced = (self.const, self._start, self._state, self._sum, self._y) #traced Arrays\n",
    "    aux_data = (self.name, self.steps)\n",
    "    # print(\"tree_flatten: \",aux_data, \" | \", children) # decomment if you want what are traced obj...\n",
    "    return traced, aux_data\n",
    "\n",
    "  @classmethod\n",
    "  def tree_unflatten(cls, aux_data, traced):\n",
    "    # print(\"tree_unflatten: \",aux_data, \" | \", children)# decomment if you want what are traced obj...\n",
    "    return cls(*aux_data, *traced)\n",
    "\n",
    "  @staticmethod\n",
    "  def _y_init(const: ArrayLike):\n",
    "    jax.debug.print('y')  # <<< print to verify it's only called once\n",
    "    return jnp.sin(const)\n",
    "\n",
    "  def _step(self, i: Array):\n",
    "    self._state += self._y\n",
    "    self._sum = self._sum.at[i].add(1)\n",
    "\n",
    "  @partial(jit, static_argnums=(1,))\n",
    "  def main(self, steps: Optional[int] = None, *, start: Array = jnp.array(0)):\n",
    "    assert self._state.ndim == 1\n",
    "    steps = steps or self.steps\n",
    "    start = start or self._start\n",
    "\n",
    "    def body(i: int, self: ExprJax):\n",
    "      jax.debug.print(\"body i:{}\",i)\n",
    "      self._step(start + i)\n",
    "      return self\n",
    "\n",
    "    # here we can do better  using scan but let keep simple\n",
    "    # The JIT will be triggered implicitly in cascade: body, _step\n",
    "    self = lax.fori_loop(0, steps, body, self)\n",
    "    self._start = start + steps\n",
    "\n",
    "    return self"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Ngja-25D3Ggv",
    "outputId": "4f9c5cd7-4425-4e52-b06b-7e526a273686"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0, 1, 2, 3, 4],\n",
       "       [5, 6, 7, 8, 9]])"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "N=10\n",
    "nsteps= 6\n",
    "starters = np.arange(N).reshape(2, -1)\n",
    "starters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "dWs6lnQFGR6g",
    "outputId": "bfe7203e-fb13-4044-f37f-c81217248988"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([-5.7535458], dtype=float32),\n",
       " array([1., 1., 1., 1., 1., 1., 0., 0., 0., 0.]))"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ExprNp('sim_np', N, np.array([5])).main(nsteps).state\n",
    "# return self._state and  self._sum where _sum is an array of N bit elements with only the first nsteps elements are 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "xuj7G1Lh4dey",
    "outputId": "aa15e269-ffb6-4cc5-a52c-a64b8c9754f0"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(array([ 0.        ,  5.048826  ,  5.4557843 ,  0.84672004, -4.540815  ],\n",
       "        dtype=float32),\n",
       "  array([1., 1., 1., 1., 1., 1., 0., 0., 0., 0.])),\n",
       " (array([-5.7535458, -1.6764929,  3.9419198,  5.9361496,  2.4727108],\n",
       "        dtype=float32),\n",
       "  array([1., 1., 1., 1., 1., 1., 0., 0., 0., 0.]))]"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#Numpy case\n",
    "res_np = [ExprNp('sim_np', N, s).main(nsteps).state for s in starters]\n",
    "res_np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "7GOB4B8R4iyh",
    "outputId": "6f8fa611-6837-4391-fe92-5136de833d92"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 0.        ,  5.04882591,  5.45578456,  0.84672005, -4.54081497],\n",
       "       [-5.75354565, -1.67649299,  3.94191959,  5.93614948,  2.47271091]])"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#just to see what is computed\n",
    "y = np.sin(starters)*nsteps\n",
    "y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "JBhm-_fxL_Kg",
    "outputId": "912eeefc-fbd6-42c0-f875-14c2fa8e2643"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[array([0, 1, 2, 3, 4]), array([5, 6, 7, 8, 9])]"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "[s for s in starters]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "mIGUJQWd6Oxg",
    "outputId": "cff21b82-4e81-4836-c5a8-fda9abe388c8"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "y\n",
      "body i:0\n",
      "body i:1\n",
      "body i:2\n",
      "body i:3\n",
      "body i:4\n",
      "body i:5\n",
      "y\n",
      "body i:0\n",
      "body i:1\n",
      "body i:2\n",
      "body i:3\n",
      "body i:4\n",
      "body i:5\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[(Array([ 0.        ,  5.0488253 ,  5.4557843 ,  0.84672004, -4.540815  ],      dtype=float32),\n",
       "  Array([1., 1., 1., 1., 1., 1., 0., 0., 0., 0.], dtype=float32)),\n",
       " (Array([-5.7535458, -1.6764929,  3.9419198,  5.9361496,  2.4727108],      dtype=float32),\n",
       "  Array([1., 1., 1., 1., 1., 1., 0., 0., 0., 0.], dtype=float32))]"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# as a list comprehension\n",
    "res_jax = [ExprJax.init('sim_jax',N, s).main(nsteps).state for s in starters]\n",
    "res_jax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "UL8s1mHTAHYm",
    "outputId": "1691c4f9-e5eb-497d-9445-b0d6152b2e14"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "y\n",
      "sim_jax:  <__main__.ExprJax object at 0x7ff4f81a10f0>\n",
      "body i:0\n",
      "body i:1\n",
      "body i:2\n",
      "body i:3\n",
      "body i:4\n",
      "body i:5\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(Array([[ 0.        ,  5.0488253 ,  5.4557843 ,  0.84672004, -4.540815  ],\n",
       "        [-5.7535458 , -1.6764929 ,  3.9419198 ,  5.9361496 ,  2.4727108 ]],      dtype=float32),\n",
       " Array([[1., 1., 1., 1., 1., 1., 0., 0., 0., 0.],\n",
       "        [1., 1., 1., 1., 1., 1., 0., 0., 0., 0.]], dtype=float32))"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# as a vmap on ExprJax objects\n",
    "sim_jax = ExprJax.init('sim_jax', N, starters)\n",
    "print(\"sim_jax: \", sim_jax)\n",
    "res_mapped = jit(vmap(partial(ExprJax.main, steps=nsteps)))(sim_jax)\n",
    "res_mapped.state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "id": "UaXNkf3wGrgb"
   },
   "outputs": [],
   "source": [
    "# A way to examine all outputs of a and b for comparison\n",
    "def tree_allclose(a, b):\n",
    "  return all(tu.tree_leaves(tu.tree_map(np.allclose, a, b)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "vvu-PZeGGzkS",
    "outputId": "e6f3662e-3798-49d6-9656-3194e258b1f4"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tree_allclose(res_jax, res_np)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "ioM2V9SEKjok",
    "outputId": "660324aa-b401-4466-f7ff-ca3cba205351"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([[ 0.        ,  5.048826  ,  5.4557843 ,  0.84672004, -4.540815  ],\n",
       "        [-5.7535458 , -1.6764929 ,  3.9419198 ,  5.9361496 ,  2.4727108 ]],\n",
       "       dtype=float32),\n",
       " array([[1., 1., 1., 1., 1., 1., 0., 0., 0., 0.],\n",
       "        [1., 1., 1., 1., 1., 1., 0., 0., 0., 0.]]))"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tu.tree_map(lambda *a: np.stack(a), *res_np)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "uQU1MTLHKs_5",
    "outputId": "ecd20988-ebc2-43e5-c8ef-39f2e3dca560"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(Array([[ 0.        ,  5.0488253 ,  5.4557843 ,  0.84672004, -4.540815  ],\n",
       "        [-5.7535458 , -1.6764929 ,  3.9419198 ,  5.9361496 ,  2.4727108 ]],      dtype=float32),\n",
       " Array([[1., 1., 1., 1., 1., 1., 0., 0., 0., 0.],\n",
       "        [1., 1., 1., 1., 1., 1., 0., 0., 0., 0.]], dtype=float32))"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "res_mapped.state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "1KYzxQEJHStr",
    "outputId": "39e19a22-5dd3-439d-cb7c-581db1fbac3a"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tree_allclose(res_mapped.state, tu.tree_map(lambda *a: np.stack(a), *res_np))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "I5VGYK-pAy2S",
    "outputId": "8f0cc72d-80b5-4308-cc4b-2ca73fc21ae6"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "y\n",
      "body i:0\n",
      "body i:1\n",
      "body i:2\n",
      "body i:0\n",
      "body i:1\n",
      "body i:2\n",
      "y\n",
      "body i:0\n",
      "body i:1\n",
      "body i:2\n",
      "body i:3\n",
      "body i:4\n",
      "body i:5\n"
     ]
    }
   ],
   "source": [
    "#using JIT & Vmap\n",
    "\n",
    "run3 = jit(vmap(partial(ExprJax.main, steps=3)))    #3 steps\n",
    "run6 = jit(vmap(partial(ExprJax.main, steps=6)))    #6 steps\n",
    "\n",
    "sim_jax33 = run3(run3(ExprJax.init('sim_jax', N, starters)))  # apply twice the 3steps function\n",
    "sim_jax6 = run6(ExprJax.init('sim_jax', N, starters))         # apply once  the 6steps function\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "gQDGCcK0Gyk1",
    "outputId": "37c7ff83-79e6-438b-bf0f-056674428672"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tree_allclose(res_jax, res_np)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "J7MXIQ1vQaCU",
    "outputId": "f3bc3d62-7cc9-491d-b79c-7fc2428f168c"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(Array([[ 0.        ,  5.0488253 ,  5.4557843 ,  0.84672004, -4.540815  ],\n",
       "        [-5.7535458 , -1.6764929 ,  3.9419198 ,  5.9361496 ,  2.4727108 ]],      dtype=float32),\n",
       " Array([[1., 1., 1., 1., 1., 1., 0., 0., 0., 0.],\n",
       "        [1., 1., 1., 1., 1., 1., 0., 0., 0., 0.]], dtype=float32))"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sim_jax33.state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "DXMkWouQQcOl",
    "outputId": "42f2440d-3c34-43d0-859f-5c2927895f62"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(Array([[ 0.        ,  5.0488253 ,  5.4557843 ,  0.84672004, -4.540815  ],\n",
       "        [-5.7535458 , -1.6764929 ,  3.9419198 ,  5.9361496 ,  2.4727108 ]],      dtype=float32),\n",
       " Array([[1., 1., 1., 1., 1., 1., 0., 0., 0., 0.],\n",
       "        [1., 1., 1., 1., 1., 1., 0., 0., 0., 0.]], dtype=float32))"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sim_jax6.state"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_XN5OdqjRwUv"
   },
   "source": [
    "# Takeaway\n",
    "\n",
    "- simple initialisation as in the case of Numpy cannot be used in JAX PyTree. Reminder: see JAX_JIT_in_class notebook to see why PyTrees should be considered.\n",
    "- initialisation should be adapted to each kind of variables later usage\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "SCepKbhHQiaM"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "authorship_tag": "ABX9TyNjhGLTIO7x8zxYYYgk/o1A",
   "include_colab_link": true,
   "provenance": []
  },
  "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
}