<commit_before>{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "M-hPMKlwXjMr"
   },
   "source": [
    "# Writing custom Jaxpr interpreters in JAX"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "r-3vMiKRYXPJ"
   },
   "source": [
    "JAX offers several composable function transformations (`jit`, `grad`, `vmap`,\n",
    "etc.) that enable writing concise, accelerated code. \n",
    "\n",
    "Here we show how to add your own function transformations to the system, by writing a custom Jaxpr interpreter. And we'll get composability with all the other transformations for free."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "s27RDKvKXFL8"
   },
   "outputs": [],
   "source": [
    "import numpy as onp\n",
    "import jax\n",
    "import jax.numpy as np\n",
    "from jax import jit, grad, vmap\n",
    "from jax import random"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "jb_8mEsJboVM"
   },
   "source": [
    "## What is JAX doing?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "KxR2WK0Ubs0R"
   },
   "source": [
    "JAX provides a NumPy-like API for numerical computing which can be used as is, but JAX's true power comes from composable function transformations. Take the `jit` function transformation, which takes in a function and returns a semantically identical function but is lazily compiled by XLA for accelerators.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {
     "height": 54
    },
    "colab_type": "code",
    "id": "HmlMcICOcSXR",
    "outputId": "546bf21b-7d03-4364-a087-5802792abbb0"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/google/home/sharadmv/workspace/jax/jax/lib/xla_bridge.py:115: UserWarning: No GPU/TPU found, falling back to CPU.\n",
      "  warnings.warn('No GPU/TPU found, falling back to CPU.')\n"
     ]
    }
   ],
   "source": [
    "x = random.normal(random.PRNGKey(0), (5000, 5000))\n",
    "def f(w, b, x):\n",
    "  return np.tanh(np.dot(x, w) + b)\n",
    "fast_f = jit(f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "gA8V51wZdsjh"
   },
   "source": [
    "When we call `fast_f`, what happens? JAX traces the function and constructs an XLA computation graph. The graph is then JIT-compiled and executed. Other transformations work similarly in that they first trace the function and handle the output trace in some way. To learn more about Jax's tracing machinery, you can refer to the [\"How it works\"](https://github.com/google/jax#how-it-works) section in the README."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "2Th1vYLVaFBz"
   },
   "source": [
    "## Jaxpr tracer\n",
    "\n",
    "A tracer of special importance in Jax is the Jaxpr tracer, which records ops into a Jaxpr (Jax expression). A Jaxpr is a data structure that can be evaluated like a mini functional programming language and \n",
    "thus Jaxprs are a useful intermediate representation\n",
    "for function transformation. \n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "pH7s63lpaHJO"
   },
   "source": [
    "To get a first look at Jaxprs, consider the `make_jaxpr` transformation. `make_jaxpr` is essentially a \"pretty-printing\" transformation:\n",
    "it transforms a function into one that, given example arguments, produces a Jaxpr representation of its computation.\n",
    "Although we can't generally use the Jaxprs that it returns, it is useful for debugging and introspection.\n",
    "Let's use it to look at how some example Jaxprs\n",
    "are structured."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {
     "height": 547
    },
    "colab_type": "code",
    "id": "RSxEiWi-EeYW",
    "outputId": "ee3d4f5c-97d3-4db3-a012-20898920cadb"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "foo\n",
      "=====\n",
      "invars: [a]\n",
      "outvars: [b]\n",
      "constvars: []\n",
      "freevars: []\n",
      "equation: [a, 1] add [b] {}\n",
      "\n",
      "jaxpr: { lambda  ;  ; a.\n",
      "  let b = add a 1\n",
      "  in [b] }\n",
      "\n",
      "\n",
      "bar\n",
      "=====\n",
      "invars: [a, b, c]\n",
      "outvars: [g, c]\n",
      "constvars: [f]\n",
      "freevars: []\n",
      "equation: [a, c] dot_general [d] {'dimension_numbers': (((1,), (0,)), ((), ())), 'precision': None}\n",
      "equation: [d, b] add [e] {}\n",
      "equation: [e, f] add [g] {}\n",
      "\n",
      "jaxpr: { lambda f ;  ; a b c.\n",
      "  let d = dot_general[ dimension_numbers=(((1,), (0,)), ((), ()))\n",
      "                       precision=None ] a c\n",
      "      e = add d b\n",
      "      g = add e f\n",
      "  in [g, c] }\n",
      "\n"
     ]
    }
   ],
   "source": [
    "def examine_jaxpr(jaxpr):\n",
    "  print(\"invars:\", jaxpr.invars)\n",
    "  print(\"outvars:\", jaxpr.outvars)\n",
    "  print(\"constvars:\", jaxpr.constvars)\n",
    "  print(\"freevars:\", jaxpr.freevars)\n",
    "  for eqn in jaxpr.eqns:\n",
    "    print(\"equation:\", eqn.invars, eqn.primitive, eqn.outvars, eqn.params)\n",
    "  print()\n",
    "  print(\"jaxpr:\", jaxpr)\n",
    "\n",
    "def foo(x):\n",
    "  return x + 1\n",
    "print(\"foo\")\n",<commit_msg>Fixed also a notebook that has gone stale<commit_after>{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "M-hPMKlwXjMr"
   },
   "source": [
    "# Writing custom Jaxpr interpreters in JAX"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "r-3vMiKRYXPJ"
   },
   "source": [
    "JAX offers several composable function transformations (`jit`, `grad`, `vmap`,\n",
    "etc.) that enable writing concise, accelerated code. \n",
    "\n",
    "Here we show how to add your own function transformations to the system, by writing a custom Jaxpr interpreter. And we'll get composability with all the other transformations for free."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "s27RDKvKXFL8"
   },
   "outputs": [],
   "source": [
    "import numpy as onp\n",
    "import jax\n",
    "import jax.numpy as np\n",
    "from jax import jit, grad, vmap\n",
    "from jax import random"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "jb_8mEsJboVM"
   },
   "source": [
    "## What is JAX doing?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "KxR2WK0Ubs0R"
   },
   "source": [
    "JAX provides a NumPy-like API for numerical computing which can be used as is, but JAX's true power comes from composable function transformations. Take the `jit` function transformation, which takes in a function and returns a semantically identical function but is lazily compiled by XLA for accelerators.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {
     "height": 54
    },
    "colab_type": "code",
    "id": "HmlMcICOcSXR",
    "outputId": "546bf21b-7d03-4364-a087-5802792abbb0"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/google/home/sharadmv/workspace/jax/jax/lib/xla_bridge.py:115: UserWarning: No GPU/TPU found, falling back to CPU.\n",
      "  warnings.warn('No GPU/TPU found, falling back to CPU.')\n"
     ]
    }
   ],
   "source": [
    "x = random.normal(random.PRNGKey(0), (5000, 5000))\n",
    "def f(w, b, x):\n",
    "  return np.tanh(np.dot(x, w) + b)\n",
    "fast_f = jit(f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "gA8V51wZdsjh"
   },
   "source": [
    "When we call `fast_f`, what happens? JAX traces the function and constructs an XLA computation graph. The graph is then JIT-compiled and executed. Other transformations work similarly in that they first trace the function and handle the output trace in some way. To learn more about Jax's tracing machinery, you can refer to the [\"How it works\"](https://github.com/google/jax#how-it-works) section in the README."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "2Th1vYLVaFBz"
   },
   "source": [
    "## Jaxpr tracer\n",
    "\n",
    "A tracer of special importance in Jax is the Jaxpr tracer, which records ops into a Jaxpr (Jax expression). A Jaxpr is a data structure that can be evaluated like a mini functional programming language and \n",
    "thus Jaxprs are a useful intermediate representation\n",
    "for function transformation. \n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "pH7s63lpaHJO"
   },
   "source": [
    "To get a first look at Jaxprs, consider the `make_jaxpr` transformation. `make_jaxpr` is essentially a \"pretty-printing\" transformation:\n",
    "it transforms a function into one that, given example arguments, produces a Jaxpr representation of its computation.\n",
    "Although we can't generally use the Jaxprs that it returns, it is useful for debugging and introspection.\n",
    "Let's use it to look at how some example Jaxprs\n",
    "are structured."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {
     "height": 547
    },
    "colab_type": "code",
    "id": "RSxEiWi-EeYW",
    "outputId": "ee3d4f5c-97d3-4db3-a012-20898920cadb"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "foo\n",
      "=====\n",
      "invars: [a]\n",
      "outvars: [b]\n",
      "constvars: []\n",
      "freevars: []\n",
      "equation: [a, 1] add [b] {}\n",
      "\n",
      "jaxpr: { lambda  ;  ; a.\n",
      "  let b = add a 1\n",
      "  in [b] }\n",
      "\n",
      "\n",
      "bar\n",
      "=====\n",
      "invars: [a, b, c]\n",
      "outvars: [g, c]\n",
      "constvars: [f]\n",
      "freevars: []\n",
      "equation: [a, c] dot_general [d] {'dimension_numbers': (((1,), (0,)), ((), ())), 'precision': None}\n",
      "equation: [d, b] add [e] {}\n",
      "equation: [e, f] add [g] {}\n",
      "\n",
      "jaxpr: { lambda f ;  ; a b c.\n",
      "  let d = dot_general[ dimension_numbers=(((1,), (0,)), ((), ()))\n",
      "                       precision=None ] a c\n",
      "      e = add d b\n",
      "      g = add e f\n",
      "  in [g, c] }\n",
      "\n"
     ]
    }
   ],
   "source": [
    "def examine_jaxpr(typed_jaxpr):\n",
    "  jaxpr = typed_jaxpr.jaxpr\n",
    "  print(\"invars:\", jaxpr.invars)\n",
    "  print(\"outvars:\", jaxpr.outvars)\n",
    "  print(\"constvars:\", jaxpr.constvars)\n",
    "  print(\"freevars:\", jaxpr.freevars)\n",
    "  for eqn in jaxpr.eqns:\n",
    "    print(\"equation:\", eqn.invars, eqn.primitive, eqn.outvars, eqn.params)\n",
    "  print()\n",
    "  print(\"jaxpr:\", jaxpr)\n",
    "\n",
    "def foo(x):\n",
    "  return x + 1\n",
    "print(\"foo\")\n",
{'fraction_non_alphanumeric': 0.1844230268299036, 'fraction_numerical': 0.01476078839975688, 'mean_word_length': 2.5212473249770713, 'pattern_counts': {'":': 140, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7317111', 'n_tokens_mistral': 4504, 'n_tokens_neox': 4307, 'n_words': 1140}