notebooks/jax/exercises/JAX_Flax_for_AI_Residents_Practical_Session_(go_flax_air) (2).ipynb

{
 "nbformat": 4,
 "nbformat_minor": 0,
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "name": "JAX/Flax for AI Residents - Practical Session (go/flax-air)",
   "provenance": [
    {
     "file_id": "/piper/depot/google3/experimental/users/andstein/jax/flax-air/practical_session.ipynb?workspaceId=andstein:flax::citc",
     "timestamp": 1605717182039
    },
    {
     "file_id": "/piper/depot/google3/experimental/users/andstein/jax/dlc/practical_session.ipynb?cl=314311380",
     "timestamp": 1605113589945
    },
    {
     "file_id": "1LHqEBdkXq3MY-6R391gxpRHoYzGkZgHh",
     "timestamp": 1590385903513
    }
   ],
   "collapsed_sections": [],
   "last_runtime": {
    "build_target": "//learning/deepmind/public/tools/ml_python:ml_notebook",
    "kind": "private"
   },
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  }
 },
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qdeIrK5I_zj_"
   },
   "source": [
    "This is the practical session of the presentation\n",
    "\"JAX/Flax for AI Residents\"\n",
    "from the AI Resident onboarding.\n",
    "\n",
    "Please see go/flax-air for **slides** and **solutions**.\n",
    "\n",
    "You probably first want to **make a copy** so you changes are not lost:\n",
    "\n",
    "![save a copy](https://screenshot.googleplex.com/34YXmEa2Lb2ipNL.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vqZcEyuRVrL4"
   },
   "source": [
    "### JAX Fundamentals"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Zdt0cskgvN9l"
   },
   "source": [
    "Now is a good moment to open the JAX documentation in a separate tab:\n",
    "\n",
    "https://jax.readthedocs.io"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "llSC6xV1Vi1O"
   },
   "source": [
    "#### Initialization"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "KGdh15C_Xh-F"
   },
   "source": [
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import numpy as np\n",
    "from matplotlib import pyplot as plt"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "1PECyFXeXg30"
   },
   "source": [
    "# Check connected accelerators. Depending on what runtime you're connected to,\n",
    "# this will show a single CPU/GPU, or 8 TPU cores (jf_2x2 aka JellyDonut).\n",
    "\n",
    "# You can start a TPU runtime via : \"Connect to a runtime\" -> \"Start\" ->\n",
    "# \"Borg Runtime\" -> \"Brain Frameworks JellyDonut (go/ml-colab)\"\n",
    "# https://screenshot.googleplex.com/87HTCpQNhBKUZUp\n",
    "# See also http://go/research-workflow-intro-deck#colab\n",
    "jax.devices()"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "fFLB_jkDYJPK"
   },
   "source": [
    "# Local devices: In this case it's the same as all devices, but if you run JAX\n",
    "# in a multi host setup, then local_devices will only show the devices connected\n",
    "# to the host running the program.\n",
    "jax.local_devices()"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "TFgHXqHBtngY"
   },
   "source": [
    "# Alternatively, you can also connect to GPU runtime.\n",
    "!nvidia-smi"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zvGr6pKaVmdf"
   },
   "source": [
    "#### Randomness\n",
    "\n",
    "https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#%F0%9F%94%AA-Random-Numbers\n",
    "\n",
    "\n",
    "<center>\n",
    "<img src=\"https://live.staticflickr.com/3127/2875827736_2224e426c6_w.jpg\" width=\"400\" height=\"300\" alt=\"Green Tree Python\"><br>\n",
    "<i>CC BY Image by <a href=\"https://www.flickr.com/photos/tedmurphy/\">Ted Murphy</a></i>\n",
    "</center>\n"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "32RCWRATVwQo"
   },
   "source": [
    "# YOUR ACTION REQUIRED:\n",
    "# Your task is to use JAX to generate 5 uniform random numbers and 5 normally\n",
    "# distributed random numbers.\n",
    "\n",
    "# Check out the following JAX API calls:\n",
    "# - jax.random.PRNGKey()\n",
    "# - jax.random.split()\n",
    "# - jax.random.uniform()\n",
    "# - jax.random.normal()"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "P5nO3BL-Vxog"
   },
   "source": [
    "#### `jnp` vs. `np`\n",
    "\n",
    "<center>\n",
    "<img src=\"https://live.staticflickr.com/2828/9578749884_d93a4a1315_w.jpg\" width=\"400\" height=\"255\" alt=\"steam forever\"><br>\n",
    "<i>CC BY Image by <a href=\"https://www.flickr.com/photos/h-studio/\">targut</a></i>\n",
    "</center>\n"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "vnix-egH3hmr"
   },
   "source": [
    "# Let's do some semi-serious matrix multiplication:\n",
    "k = 3_000\n",
    "x = np.random.normal(size=[k, k])\n",
    "# ~3.4s\n",
    "%time x @ x"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "6JwTdcQH3Vhp"
   },
   "source": [
    "# YOUR ACTION REQUIRED: Do the same computation using JAX!\n",
    "# You should use result.block_until_ready() for a fair comparison."
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "aAfR3DNTNRG0"
   },
   "source": [
    "# Note the different class of the JAX array. There is additional API e.g. to\n",
    "# determine on which device the data is stored, check out x.device_buffer"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "8L1Vzbzj234K"
   },
   "source": [
    "# Combining jnp & np : Below array initialization is rather slow because we\n",
    "# create a lot of jnp array. Replace jnp with np and observe the speedup!\n",
    "%%time\n",
    "# GPU : 1.79s\n",
    "# CPU : 1.04s\n",
    "x = jnp.array([jnp.arange(100) for _ in range(10000)])\n",
    "print(repr(x))\n",
    "# YOUR ACTION REQUIRED:\n",
    "# In this situation we would want to create the array in np and then convert it\n",
    "# to a jnp array using jnp.array() or jax.device_put().\n",
    "# (Note that we could use np.tile() here, but that's not the point)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "I3U44rDeVo--"
   },
   "source": [
    "#### `grad()`\n",
    "\n",
    "<center>\n",
    "<img src=\"https://live.staticflickr.com/8573/15246394073_0cfdcc458b_w.jpg\" width=\"400\" height=\"221\" alt=\"Gradient\"><br>\n",
    "<i>CC BY Image by <a href=\"https://www.flickr.com/photos/60506610@N08/\">Manel Torralba</a></i>\n",
    "</center>\n",
    "\n",
    "https://jax.readthedocs.io/en/latest/jax.html#automatic-differentiation"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "OmqVcxf8W0PL"
   },
   "source": [
    "def sigmoid(x):\n",
    "    return 0.5 * (1 + jnp.tanh(x))\n",
    "\n",
    "\n",
    "# YOUR ACTION REQUIRED:\n",
    "# Use grad() to create a new function that computes the gradient of `sigmoid`.\n",
    "# Verify the output of the new function at some points."
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "vttdLwiCxfFl"
   },
   "source": [
    "def f(x, y):\n",
    "    return 2 * x * y**2\n",
    "\n",
    "\n",
    "# YOUR ACTION REQUIRED:\n",
    "# Compute df/dx and df/dy with grad()"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KsVY8VxEV5HK"
   },
   "source": [
    "#### `vmap()`\n",
    "\n",
    "<center>\n",
    "<img src=\"https://live.staticflickr.com/65535/49164406707_a954dc465f_w.jpg\" width=\"400\" height=\"225\" alt=\"Les Tanji, éléments majeurs du paysage urbain coréen (Daejeon, Corée du sud)\"><br>\n",
    "<i>CC BY Image by <a href=\"https://www.flickr.com/photos/dalbera/\">Jean-Pierre Dalbéra</a></i>\n",
    "</center>\n",
    "\n",
    "https://jax.readthedocs.io/en/latest/jax.html#vectorization-vmap"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "4wsPRBn1W0gq"
   },
   "source": [
    "# Now let's plot the gradient of the sigmoid function in the range [05, 5]\n",
    "xs = jnp.linspace(-5, 5, 100)\n",
    "# We can of course evaluate the gradient at every position separately:\n",
    "grads = [jax.grad(sigmoid)(x) for x in xs]\n",
    "plt.plot(xs, grads);\n",
    "# But JAX can \"vectorize\" our gradient function for us automatically.\n",
    "# YOUR ACTION REQUIRED:\n",
    "# Read the documentation about `vmap` and reimplement the plot without a Python\n",
    "# loop."
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "AWtpfcOy01lh"
   },
   "source": [
    "# Another vmap() example : Let's re-implement matmul using vector dot product:\n",
    "vdp = lambda v1, v2: v1.dot(v2)\n",
    "# Vector dot product:\n",
    "vdp(jnp.arange(1, 4), jnp.arange(1, 4))\n",
    "# Matrix vector product:\n",
    "mvp = jax.vmap(vdp, in_axes=(0, None), out_axes=0)\n",
    "# Matrix matrix product:\n",
    "mmp = jax.vmap(mvp, in_axes=(None, 1), out_axes=1)\n",
    "\n",
    "# Verify result.\n",
    "m1 = jnp.arange(12).reshape((3, 4))\n",
    "m2 = m1.reshape((4, 3))\n",
    "# In case you were wondering : Since Python 3.5 we have `.__matmul__()` operator\n",
    "# that happens to use the same character as for decorators (cf. `@jit` below).\n",
    "mmp(m1, m2) - m1 @ m2\n",
    "\n",
    "# YOUR ACTION REQUIRED:\n",
    "# It's curry time!\n",
    "# Try re-implementing mvp() but this time without using the in_axes=, and\n",
    "# out_axes=. Instead use lambda expressions to (un)curry the arguments in such\n",
    "# a way that vmap()'s default in_axes=0 and out_axes=0 does the job.\n",
    "# (You can also re-implement mmp() this way, but it involves transposing)."
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "sNoQ-NVBV7nM"
   },
   "source": [
    "#### `jit()`\n",
    "\n",
    "<center>\n",
    "<img src=\"https://live.staticflickr.com/3803/9540184355_0dee2f496a_w.jpg\" width=\"400\" height=\"267\" alt=\"Silicon Village\"><br>\n",
    "<i>CC BY Image by <a href=\"https://www.flickr.com/photos/jackofspades/\">Jack Spades</a></i>\n",
    "</center>\n",
    "\n",
    "https://jax.readthedocs.io/en/latest/jax.html#just-in-time-compilation-*jit*"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "R_3LkmBs8tjh"
   },
   "source": [
    "# JAX would not have the final X in it's name if it were not for XLA, the\n",
    "# magic sauce that somehow takes computation defined in a function as input\n",
    "# and produces a much faster version of it.\n",
    "\n",
    "# @jax.jit\n",
    "def f(x):\n",
    "    y = x\n",
    "    for _ in range(10):\n",
    "        y = y - 0.1 * y + 3.0\n",
    "    return y[:100, :100]\n",
    "\n",
    "\n",
    "x = jax.random.normal(jax.random.PRNGKey(0), (5000, 5000))\n",
    "%timeit f(x).block_until_ready()\n",
    "\n",
    "# YOUR ACTION REQUIRED:\n",
    "# Move your magic JAX wand and cast a spell by removing a single character from\n",
    "# above example, drastically speeding up the computation!\n",
    "# Note: JIT unrolls the for loop and converts all computations to XLA\n",
    "# primitives. XLA is then smart enough to fuse kernels for multiplication and\n",
    "# addition, and optimize the program to only compute those parts that are\n",
    "# actually needed for the function result..."
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "cKN0v0sCEZqq"
   },
   "source": [
    "# Just to be clear : `@jit` is Python's decorator syntax [1], you can also use\n",
    "# jit() like the other function transformations.\n",
    "# [1] https://www.python.org/dev/peps/pep-0318\n",
    "\n",
    "\n",
    "@jax.jit\n",
    "def f1_jit(x):\n",
    "    return x**0.5\n",
    "\n",
    "\n",
    "def f2(x):\n",
    "    return x**0.5\n",
    "\n",
    "\n",
    "# It's really the same.\n",
    "f2_jit = jax.jit(f2)\n",
    "\n",
    "f1_jit(2) - f2_jit(2)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "B2G-A2tnvqao"
   },
   "source": [
    "# What you need to understand about JIT (1/3): When a function is traced.\n",
    "\n",
    "\n",
    "@jax.jit\n",
    "def noop(x):\n",
    "    # This statement only gets executed when the function is traced, i.e. every\n",
    "    # time you execute the JIT-ted version with a new ShapedArray (different dtype\n",
    "    # and/or different shape).\n",
    "    print(\"Tracing noop:\", x)\n",
    "    return x\n",
    "\n",
    "\n",
    "noop(jnp.arange(3))  # Tracing.\n",
    "noop(jnp.arange(3) + 1)  # Using trace from cache.\n",
    "noop(jnp.arange(4))  # Tracing.\n",
    "noop(jnp.arange(4.0))  # Tracing.\n",
    "noop(jnp.arange(1.0, 5.0))  # Using trace from cache."
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "9gKO9M7DDJUr"
   },
   "source": [
    "# What you need to understand about JIT (2/3): Baking in environment.\n",
    "magic_number = 13\n",
    "\n",
    "\n",
    "@jax.jit\n",
    "def add_magic(x):\n",
    "    return x + magic_number\n",
    "\n",
    "\n",
    "print(add_magic(np.array([0])))\n",
    "magic_number = 42\n",
    "print(add_magic(np.array([0])))\n",
    "print(add_magic(np.array([0.0])))"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "6neI1cN8Dgho"
   },
   "source": [
    "# What you need to understand about JIT (2/3): Value-dependent flow.\n",
    "def mult(x, n):\n",
    "    print(\"Tracing mult:\", x, n)\n",
    "    tot = 0\n",
    "    while n > 0:\n",
    "        tot += x\n",
    "        n -= 1\n",
    "    return tot"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "kudcEghERJF5"
   },
   "source": [
    "# The problem:\n",
    "\n",
    "# The following statement fails, because : JIT will generate the function's XLA\n",
    "# code by tracing it with `ShapedArray`'s. These arrays have only their shape\n",
    "# and datatype defined. Hence, if there are any statements involving the actual\n",
    "# *values* of the parameters, JIT does not know what to do and raises an\n",
    "# exception.\n",
    "# (Note that if mult were traced with `ConcreteArray`s then the trace would work\n",
    "#  just fine; you can see that when executing `grad(mult)(3., 2.)`)\n",
    "try:\n",
    "    jax.jit(mult)(3, 2)\n",
    "except Exception as e:\n",
    "    print(f\"\\n### FAILED WITH : {e}\")\n",
    "\n",
    "# How can we fix this ??"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "BDMWMsPkGToO"
   },
   "source": [
    "# Solution 1 : static_argnums\n",
    "jax.jit(mult, static_argnums=1)(3, 4)\n",
    "jax.jit(mult, static_argnums=1)(3, 5)\n",
    "jax.jit(mult, static_argnums=1)(3, 6)\n",
    "\n",
    "# By the way : did you notice how the function is traced exactly three times the\n",
    "# first time this cell is executed, but not when you re-execute the same cell?\n",
    "# That's because JIT-ted functions are cached. If You want to observe the\n",
    "# tracing a second time, you first need to execute above cell so that `mult`\n",
    "# gets redefined and the cache needs to be updated with the new definition."
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "PVg2Df7IHyro"
   },
   "source": [
    "# Solution 2 : (un)currying\n",
    "\n",
    "# YOUR ACTION REQUIRED:\n",
    "# Use jit() without `static_argnums=`, but (un)curry the function mult instead."
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "BcLw1EuyIHeV"
   },
   "source": [
    "# Solution 3 : Use XLA primitives for control flow.\n",
    "\n",
    "# Remember: You can inspect `jax.lax.while_loop()` docs by either:\n",
    "# - Go to https://jax.readthedocs.io\n",
    "# - Execute a cell containing `?jax.lax.while_loop`\n",
    "# - Hover your mouse over `while_loop` and wait two seconds\n",
    "\n",
    "\n",
    "def mult_(x, n):\n",
    "    print(\"Tracing mult_:\", x, n)\n",
    "\n",
    "    def cond_fun(n_tot):\n",
    "        n, tot = n_tot\n",
    "        return n > 0\n",
    "\n",
    "    def body_fun(n_tot):\n",
    "        n, tot = n_tot\n",
    "        return (n - 1, tot + x)\n",
    "\n",
    "    return jax.lax.while_loop(cond_fun, body_fun, (n, 0))\n",
    "\n",
    "\n",
    "jax.jit(mult_)(3, 4)\n",
    "jax.jit(mult_)(3, 5)\n",
    "jax.jit(mult_)(3, 6)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "4DIdnlULSoDs"
   },
   "source": [
    "# Woah! Wasn't JAX supposed to be fast !? What is going on here ??\n",
    "# Also note that increasing the second number significantly will crash\n",
    "# your runtime...\n",
    "%%time\n",
    "jax.jit(mult, static_argnums=1)(3, 5000)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "ud2GC4vhTJDR"
   },
   "source": [
    "# Does this function have the same problems? Why not?\n",
    "%%time\n",
    "jax.jit(mult_)(3, 5000)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qS2kXGXFG9I-"
   },
   "source": [
    "#### `pmap()`\n",
    "\n",
    "<center>\n",
    "<a href=\"https://storage.googleapis.com/gweb-cloudblog-publish/original_images/TPU_V3_POD_FULLFRONT_FORWEBONLY_FINAL.jpg\" target=\"_blank\"><img src=\"https://storage.googleapis.com/gweb-cloudblog-publish/original_images/TPU_V3_POD_FULLFRONT_FORWEBONLY_FINAL.jpg\" width=\"800\"  alt=\"Full TPUv3 (DragonFish) pod\"></a><br>\n",
    "<i>Full TPUv3 (DragonFish) pod (from <a href=\"https://cloud.google.com/blog/products/ai-machine-learning/cloud-tpu-pods-break-ai-training-records?hl=fr_ca&skip_cache=true\">GoogleAI Blog</a>) - click to enlarge.</i>\n",
    "</center>\n",
    "\n",
    "https://jax.readthedocs.io/en/latest/jax.html#parallelization-pmap"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "XQ0344DG1w3T"
   },
   "source": [
    "# Parallel computing is more fun with multiple devices :-)\n",
    "# Go back to \"Initialization\" and connect to a different runtime if you're\n",
    "# running on a single device.\n",
    "assert jax.device_count() == 8, \"Please connect to a JellyDonut runtime!\""
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "0cfM4Cl42V3K"
   },
   "source": [
    "# By default in_axes=0, so pmap() will split every incoming tensor across it's\n",
    "# first axis - which should be sized jax.local_device_count().\n",
    "# The computations are then performed in parallel and the results are returned\n",
    "# as a sharded device array. The dat remains on the individual accelerators.\n",
    "# Note that pmap() also XLA-compiles the function, so no need to call jit().\n",
    "\n",
    "# Generate 8 different random seeds.\n",
    "keys = jax.random.split(jax.random.PRNGKey(0), 8)\n",
    "# Generate 8 different random matrices. Data remains on devices.\n",
    "mats = jax.pmap(lambda key: jax.random.normal(key, (8_000, 8_000)))(keys)\n",
    "# Perform 8 matmuls in parallel.\n",
    "results = jax.pmap(lambda m1, m2: m1 @ m2)(mats, mats)\n",
    "\n",
    "# YOUR ACTION REQUIRED:\n",
    "# Fetch the mean of thes matrices from every device and print it out here."
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "SI0FUo5jyGiD"
   },
   "source": [
    "import functools\n",
    "\n",
    "# Here we use jax.lax.psum() to do computations across devices. Note that these\n",
    "# operations can cause a lot of communication costs. Below we split our 8\n",
    "# devices along two axis (4x2).\n",
    "\n",
    "# Note in particular that parallel operators work across hosts! We can't\n",
    "# demonstrate this in a Colab, but you will encounter it later in the Flax\n",
    "# examples and brain templates.\n",
    "\n",
    "# You can read more about parallel operators here:\n",
    "# https://jax.readthedocs.io/en/latest/jax.lax.html#parallel-operators\n",
    "\n",
    "# axis 0 : rows\n",
    "@functools.partial(jax.pmap, axis_name=\"rows\")\n",
    "# axis 1 : columns\n",
    "@functools.partial(jax.pmap, axis_name=\"cols\")\n",
    "def f(x):\n",
    "    # across the rows (= column sum)\n",
    "    row_sum = jax.lax.psum(x, \"rows\")\n",
    "    # across the cols (= row sum)\n",
    "    col_sum = jax.lax.psum(x, \"cols\")\n",
    "    total_sum = jax.lax.psum(x, (\"rows\", \"cols\"))\n",
    "    return row_sum, col_sum, total_sum\n",
    "\n",
    "\n",
    "# YOUR ACTION REQUIRED:\n",
    "# Create an array, feed it to f() and verify the correctness of the results"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "rkgsBV1nbQRC"
   },
   "source": [
    "#### pytrees\n",
    "\n",
    "<center>\n",
    "<img src=\"https://live.staticflickr.com/4695/38641518410_53da16c2a9_w.jpg\" width=\"400\" height=\"300\" alt=\"Green Tree Python\"><br>\n",
    "<i>CC BY Image by <a href=\"https://www.flickr.com/photos/markgillow/\">Mark Gillow</a></i>\n",
    "</center>\n",
    "\n",
    "https://jax.readthedocs.io/en/latest/notebooks/JAX_pytrees.html"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "bsJGdklgbRMo"
   },
   "source": [
    "# Whenever we encounter a function argument, e.g. the `params` for a model, or\n",
    "# the first argument to `grad()` to whose respect we perform automatic\n",
    "# differentiation, it can really be a \"pytree\" of `jnp.ndarray`. A pytree\n",
    "# consists of an arbitrary combination of Python dict/list/tuple and allows us\n",
    "# to structure our data hierarchically.\n",
    "\n",
    "# This is a pytree:\n",
    "data = dict(\n",
    "    array_3x2=jnp.arange(6.0).reshape((3, 2)),\n",
    "    mixed_tuple=(0.1, 0.2, 0.3, [1.0, 2.0, 3.0]),\n",
    "    subdict=dict(\n",
    "        array_3x4=jnp.arange(12.0).reshape((3, 4)),\n",
    "        array_4x3=jnp.arange(12.0).reshape((4, 3)),\n",
    "    ),\n",
    ")"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "yapH9c-teEhY"
   },
   "source": [
    "# Call a function over all values, output resulting tree:\n",
    "jax.tree_map(jnp.shape, data)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "CycIGKXWeHh7"
   },
   "source": [
    "# Define a function that does some computation with the values:\n",
    "def sumsquares(x):\n",
    "    value_flat, value_tree = jax.tree_flatten(x)\n",
    "    del value_tree  # not needed.\n",
    "    tot = 0\n",
    "    for value in value_flat:\n",
    "        if isinstance(value, jnp.ndarray):\n",
    "            value = value.sum()\n",
    "        tot += value**2\n",
    "    return tot\n",
    "\n",
    "\n",
    "sumsquares(data)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "0jlfIFXIejG4"
   },
   "source": [
    "# Compute gradients. Remember that grad() computes gradients wrt the first\n",
    "# argument, but that first argument can be an arbitrarily complex pytree (like\n",
    "# all the weights in your hierarchical model).\n",
    "grads = jax.grad(sumsquares)(data)\n",
    "grads"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "yr0j_wYlfncp"
   },
   "source": [
    "# YOUR ACTION REQUIRED:\n",
    "# Take a step against the gradients using `jax.tree_multimap()`"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "f6k6zea5VwZ_"
   },
   "source": [
    "### JAX Linear Classifier"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "bWcGWdZVpLBz"
   },
   "source": [
    "#### Fetch data"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "_4a6locEVxif"
   },
   "source": [
    "# Our one stop shop for datasets. If you use dataset preprocessing, then those\n",
    "# computations will be performed with a Tensorflow graph. But we don't really\n",
    "# need to understand the details, but rather use the API to stream through the\n",
    "# dataset and then use JAX for computations.\n",
    "import tensorflow_datasets as tfds"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "UZblqEC2Wgwy"
   },
   "source": [
    "# Don't like fashion? Go checkout the other image classification datasets:\n",
    "# https://www.tensorflow.org/datasets/catalog/overview#image_classification\n",
    "# (actually, go and check them out, even if you like fashion...)\n",
    "ds, ds_info = tfds.load(\"fashion_mnist\", with_info=True)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "RJIzkjNIW1Jh"
   },
   "source": [
    "tfds.show_examples(ds[\"train\"], ds_info, rows=4, cols=6);"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "bKcAy-HXJ3I5"
   },
   "source": [
    "# We're not really interested in tf.data preprocessing here, so let's just fetch\n",
    "# all the data as a jax.ndarray...\n",
    "\n",
    "\n",
    "def ds_get_all(ds, *keys):\n",
    "    \"\"\"Returns jnp.array() for specified `keys` from entire dataset `ds`.\"\"\"\n",
    "    d = next(iter(ds.batch(ds.cardinality())))\n",
    "    return tuple(jnp.array(d[key]._numpy()) for key in keys)\n",
    "\n",
    "\n",
    "train_images, train_labels = ds_get_all(ds[\"train\"], \"image\", \"label\")\n",
    "train_images /= 255.0\n",
    "test_images, test_labels = ds_get_all(ds[\"test\"], \"image\", \"label\")\n",
    "test_images /= 255.0\n",
    "\n",
    "train_images.shape, train_labels.shape  # labels as indices, not one-hot"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9YBzuATupOAS"
   },
   "source": [
    "#### Step 1 : Define a model"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "SmoTK4nTZOl2"
   },
   "source": [
    "# YOUR ACTION REQUIRED:\n",
    "# Implement the body of this function.\n",
    "def linear_init(key, input_shape, n_classes):\n",
    "    \"\"\"Initializes parameters for a linear classifier.\n",
    "\n",
    "    Args:\n",
    "      key: a PRNGKey used as the random key.\n",
    "      input_shape: Shape of a single input example.\n",
    "      n_classes: Number of output classes.\n",
    "\n",
    "    Returns:\n",
    "      A pytree to be used as a first argument with `linear_apply()`.\n",
    "    \"\"\"\n",
    "    pass\n",
    "\n",
    "\n",
    "# YOUR ACTION REQUIRED:\n",
    "# Implement the body of this function.\n",
    "def linear_apply(params, inp):\n",
    "    \"\"\"Computes logits for a SINGLE EXAMPLE.\n",
    "\n",
    "    Args:\n",
    "      params: A pytree as returned by `linear_init()`.\n",
    "      inp: A single input example.\n",
    "\n",
    "    Returns:\n",
    "      Logits (i.e. values that should be normalized by `jax.nn.softmax()` to get a\n",
    "      valid probability distribution over the output classes).\n",
    "    \"\"\"\n",
    "    pass"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "iPuXRp4QhC4-"
   },
   "source": [
    "# Initialize classifier & run on a single example.\n",
    "\n",
    "params = linear_init(\n",
    "    key=jax.random.PRNGKey(0),\n",
    "    input_shape=train_images[0].shape,\n",
    "    n_classes=ds_info.features[\"label\"].num_classes,\n",
    ")\n",
    "print(jax.tree_map(jnp.shape, params))\n",
    "linear_apply(params, train_images[0])"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "tk6o3ypvpSGj"
   },
   "source": [
    "#### Step 2 : Define a loss"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "OMQk7whohXd7"
   },
   "source": [
    "def loss_fun(params, inputs, targets):\n",
    "    \"\"\"Compute x-entropy loss for a batch of images.\n",
    "\n",
    "    Args:\n",
    "      params: a pytree as returned by `linear_init()`.\n",
    "      inputs: batch of images\n",
    "      targets: batch of target labels (indices)\n",
    "\n",
    "    Returns:\n",
    "      The loss value.\n",
    "    \"\"\"\n",
    "    # Note that we defined linear_apply() for a single example and how we use\n",
    "    # `vmap()` here to vectorize the function.\n",
    "    logits = jax.vmap(linear_apply, in_axes=(None, 0))(params, inputs)\n",
    "    # We go from logits directly to log(probs):\n",
    "    logprobs = logits - jax.scipy.special.logsumexp(\n",
    "        logits, axis=-1, keepdims=True\n",
    "    )\n",
    "    # Note: targets are indices.\n",
    "    return -logprobs[jnp.arange(len(targets)), targets].mean()\n",
    "\n",
    "\n",
    "loss_fun(params, train_images[:2], train_labels[:2])"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "pNgqVkuVpVZE"
   },
   "source": [
    "#### Step 3 : `update_step()`"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "EaXX-bHlhYjE"
   },
   "source": [
    "# This is a good moment to compile our computations using `jit()` !\n",
    "# REMEMBER: Since we \"bake in\" all globals when `jit()` is called, you will need\n",
    "# to re-execute this cell every time you change some code `update_step()`\n",
    "# depends on (like e.g. `loss_fun()`, or `linear_apply()`).\n",
    "@jax.jit\n",
    "def update_step(params, inputs, targets):\n",
    "    \"\"\"Take a single optimization step.\n",
    "\n",
    "    Args:\n",
    "      params: A pytree as returned by `linear_init()`.\n",
    "      inputs: batch of images\n",
    "      targets: batch of target labels (indices)\n",
    "\n",
    "    Returns:\n",
    "      A tuple (updated_params, loss).\n",
    "    \"\"\"\n",
    "    loss, grads = jax.value_and_grad(loss_fun)(params, inputs, targets)\n",
    "    # Opimize using SGD\n",
    "    updated_params = jax.tree_multimap(\n",
    "        lambda param, grad: param - 0.05 * grad, params, grads\n",
    "    )\n",
    "    return updated_params, loss\n",
    "\n",
    "\n",
    "update_step(params, train_images[:2], train_labels[:2])"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "txVsdtoWpcdJ"
   },
   "source": [
    "#### Step 4: Train the model"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "K4q9JHQLhelD"
   },
   "source": [
    "# Step 4 : Do the training by calling `update_step()` repeatedly.\n",
    "\n",
    "\n",
    "def train(params, steps, batch_size=128):\n",
    "    losses = []\n",
    "    steps_per_epoch = len(train_images) // batch_size\n",
    "    for step in range(steps):\n",
    "        i0 = (step % steps_per_epoch) * batch_size\n",
    "        # Training is simply done by calling `update_step()` repeatedly and\n",
    "        # replacing `params` with `updated_params` returned by `update_step()`.\n",
    "        params, loss = update_step(\n",
    "            params,\n",
    "            train_images[i0 : i0 + batch_size],\n",
    "            train_labels[i0 : i0 + batch_size],\n",
    "        )\n",
    "        losses.append(float(loss))\n",
    "    return params, jnp.array(losses)\n",
    "\n",
    "\n",
    "learnt_params, losses = train(params, steps=1_000)\n",
    "plt.plot(losses)\n",
    "print(\"final loss:\", np.mean(losses[-100]))"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "S_1_9-ISoGyd"
   },
   "source": [
    "# Compute accuracy of linear model.\n",
    "\n",
    "\n",
    "def accuracy(params, inputs, targets):\n",
    "    logits = jax.vmap(linear_apply, in_axes=(None, 0))(params, inputs)\n",
    "    return (targets == logits.argmax(axis=-1)).mean()\n",
    "\n",
    "\n",
    "accuracy(learnt_params, test_images, test_labels)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "N8Vf0fZqPBY4"
   },
   "source": [
    "### Flax\n",
    "\n",
    "You probably want to keep the Flax documentation ready in another tab:\n",
    "\n",
    "https://flax.readthedocs.io/"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "P7YKHVVoVoQS"
   },
   "source": [
    "# from typing import Callable, Sequence  # used ?\n",
    "\n",
    "import flax\n",
    "from flax import linen as nn"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "uiR7qpQy9ksY"
   },
   "source": [
    "#### Functional core"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "E4qwAAEVE_hp"
   },
   "source": [
    "# Simple module with matmul layer. Note that we could build this in many\n",
    "# different ways using the `scope` for parameter handling.\n",
    "\n",
    "\n",
    "class Matmul:\n",
    "    def __init__(self, features):\n",
    "        self.features = features\n",
    "\n",
    "    def kernel_init(self, key, shape):\n",
    "        return jax.random.normal(key, shape)\n",
    "\n",
    "    def __call__(self, scope, x):\n",
    "        kernel = scope.param(\n",
    "            \"kernel\", self.kernel_init, (x.shape[1], self.features)\n",
    "        )\n",
    "        return x @ kernel\n",
    "\n",
    "\n",
    "class Model:\n",
    "    def __init__(self, features):\n",
    "        self.matmuls = [Matmul(f) for f in features]\n",
    "\n",
    "    def __call__(self, scope, x):\n",
    "        x = x.reshape([len(x), -1])\n",
    "        for i, matmul in enumerate(self.matmuls):\n",
    "            x = scope.child(matmul, f\"matmul_{i + 1}\")(x)\n",
    "            if i < len(self.matmuls) - 1:\n",
    "                x = jax.nn.relu(x)\n",
    "        x = jax.nn.log_softmax(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "model = Model([ds_info.features[\"label\"].num_classes])\n",
    "y, variables = flax.core.init(model)(key, train_images[:1])\n",
    "assert (y == flax.core.apply(model)(variables, train_images[:1])).all()\n",
    "\n",
    "# YOUR ACTION REQUIRED:\n",
    "# Check out the parameter structure, try adding/removing \"layers\" and see how it\n",
    "# changes"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "dxEq47C2Fite"
   },
   "source": [
    "# YOUR ACTION REQUIRED:\n",
    "# Redefine loss_fun(), update_step(), and train() from above to train the new\n",
    "# model."
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2nJSoLKUCbnF"
   },
   "source": [
    "#### Stateless Linen module"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "1ciRbfVPRvb9"
   },
   "source": [
    "# Reimplementation of above model using the Linen API.\n",
    "\n",
    "\n",
    "class Model(nn.Module):\n",
    "    num_classes: int\n",
    "\n",
    "    def setup(self):\n",
    "        self.dense = nn.Dense(self.num_classes)\n",
    "\n",
    "    def __call__(self, x):\n",
    "        x = x.reshape([len(x), -1])\n",
    "        x = self.dense(x)\n",
    "        x = nn.log_softmax(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "model = Model(num_classes=ds_info.features[\"label\"].num_classes)\n",
    "variables = model.init(jax.random.PRNGKey(0), train_images[:1])\n",
    "jax.tree_map(jnp.shape, variables)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "hyNbZwFS8vlI"
   },
   "source": [
    "# YOUR ACTION REQUIRED:\n",
    "# 1. Rewrite above model using the @nn.compact notation.\n",
    "# 2. Extend the model to use additional layers, see e.g.\n",
    "#    convolutions in\n",
    "#    http://google3/third_party/py/flax/linen/linear.py"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "dah65MVi-bRR"
   },
   "source": [
    "model = Model(ds_info.features[\"label\"].num_classes)\n",
    "variables = model.init(key, train_images[:1])\n",
    "jax.tree_map(jnp.shape, variables)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "dOCSpW5sR3fk"
   },
   "source": [
    "# Reimplementation of training loop using a Flax optimizer.\n",
    "\n",
    "\n",
    "@jax.jit\n",
    "def update_step_optim(optim, inputs, targets):\n",
    "    def loss_fun(params):\n",
    "        logits = model.apply(dict(params=params), inputs)\n",
    "        logprobs = logits - jax.scipy.special.logsumexp(\n",
    "            logits, axis=-1, keepdims=True\n",
    "        )\n",
    "        return -logprobs[jnp.arange(len(targets)), targets].mean()\n",
    "\n",
    "    loss, grads = jax.value_and_grad(loss_fun)(optim.target)\n",
    "    return optim.apply_gradient(grads), loss\n",
    "\n",
    "\n",
    "def train_optim(optim, steps, batch_size=128):\n",
    "    losses = []\n",
    "    steps_per_epoch = len(train_images) // batch_size\n",
    "    for step in range(steps):\n",
    "        i0 = (step % steps_per_epoch) * batch_size\n",
    "        optim, loss = update_step_optim(\n",
    "            optim,\n",
    "            train_images[i0 : i0 + batch_size],\n",
    "            train_labels[i0 : i0 + batch_size],\n",
    "        )\n",
    "        losses.append(float(loss))\n",
    "    return optim, jnp.array(losses)\n",
    "\n",
    "\n",
    "optim = flax.optim.adam.Adam(learning_rate=0.01).create(variables[\"params\"])\n",
    "learnt_optim, losses = train_optim(optim, steps=1_000)\n",
    "plt.plot(losses)\n",
    "print(\"final loss:\", np.mean(losses[-100]))"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "7S2UhjE-SLfi"
   },
   "source": [
    "# Re-evaluate accuracy.\n",
    "(\n",
    "    model.apply(dict(params=learnt_optim.target), test_images).argmax(axis=-1)\n",
    "    == test_labels\n",
    ").mean()"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "fMm28tSqWRkc"
   },
   "source": [
    "#### Linen module with state"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "YfsyKCUSezL5"
   },
   "source": [
    "# Let's add batch norm!\n",
    "# I'm not saying it's a good idea here, but it will allow us study the changes\n",
    "# we need to make for models that have state.\n",
    "\n",
    "\n",
    "class Model(nn.Module):\n",
    "    num_classes: int\n",
    "\n",
    "    @nn.compact\n",
    "    def __call__(self, x, *, train):\n",
    "        x = x.reshape([len(x), -1])\n",
    "        x = nn.BatchNorm(use_running_average=not train)(x)\n",
    "        x = nn.Dense(self.num_classes)(x)\n",
    "        x = nn.log_softmax(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "model = Model(num_classes=ds_info.features[\"label\"].num_classes)\n",
    "variables = model.init(jax.random.PRNGKey(0), train_images[:1], train=True)\n",
    "jax.tree_map(jnp.shape, variables)\n",
    "\n",
    "# Note the new \"batch_stats\" collection !"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "EuF11BFGCku7"
   },
   "source": [
    "# YOUR ACTION REQUIRED:\n",
    "# Check below code and add comments for every change compared to the model above\n",
    "# without state.\n",
    "\n",
    "\n",
    "@jax.jit\n",
    "def update_step_optim(optim, batch_stats, inputs, targets):\n",
    "    def loss_fun(params):\n",
    "        logits, mutated_state = model.apply(\n",
    "            dict(params=params, batch_stats=batch_stats),\n",
    "            inputs,\n",
    "            mutable=\"batch_stats\",\n",
    "            train=True,\n",
    "        )\n",
    "        logprobs = logits - jax.scipy.special.logsumexp(\n",
    "            logits, axis=-1, keepdims=True\n",
    "        )\n",
    "        return (\n",
    "            -logprobs[jnp.arange(len(targets)), targets].mean(),\n",
    "            variables[\"batch_stats\"],\n",
    "        )\n",
    "\n",
    "    (loss, state), grads = jax.value_and_grad(loss_fun, has_aux=True)(\n",
    "        optim.target\n",
    "    )\n",
    "    return optim.apply_gradient(grads), batch_stats, loss\n",
    "\n",
    "\n",
    "def train_optim(optim, batch_stats, steps, batch_size=128):\n",
    "    losses = []\n",
    "    steps_per_epoch = len(train_images) // batch_size\n",
    "    for step in range(steps):\n",
    "        i0 = (step % steps_per_epoch) * batch_size\n",
    "        optim, batch_stats, loss = update_step_optim(\n",
    "            optim,\n",
    "            batch_stats,\n",
    "            train_images[i0 : i0 + batch_size],\n",
    "            train_labels[i0 : i0 + batch_size],\n",
    "        )\n",
    "        losses.append(float(loss))\n",
    "    return optim, batch_stats, jnp.array(losses)\n",
    "\n",
    "\n",
    "optim = flax.optim.adam.Adam(learning_rate=0.01).create(variables[\"params\"])\n",
    "learnt_optim, batch_stats, losses = train_optim(\n",
    "    optim, variables[\"batch_stats\"], steps=1_000\n",
    ")\n",
    "plt.plot(losses)\n",
    "print(\"final loss:\", np.mean(losses[-100]))"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "Vk8-FFi3EN-Q"
   },
   "source": [
    "# YOUR ACTION REQUIRED:\n",
    "# Make predictions with above model with state"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "CczOcHDOe-34"
   },
   "source": [
    "#### Modify MNIST example\n",
    "\n",
    "Check out the Flax MNIST example Colab - you can find a link on Github\n",
    "\n",
    "https://github.com/google/flax/tree/master/linen_examples/mnist"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "8PLQpjaJfL40"
   },
   "source": [
    "# YOUR ACTION REQURIED:\n",
    "# Store the Colab in your personal drive and modify it to use the dataset from\n",
    "# above.\n",
    "# While this might sound boring, you will learn the following things:\n",
    "# - how to load files in public Colab from Github, modify them in the UI and\n",
    "#   optionally store them on your personal Google Drive.\n",
    "# - how to use inline TensorBoard on public Colab and export it to tensorboard.dev"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "drWpnWgcO_hp"
   },
   "source": [
    "### Brain templates\n",
    "\n",
    "code : go/brain-templates\n",
    "\n",
    "documentation : go/brain-templates-doc\n",
    "\n",
    "These are more open ended exercises, but they could well pay off most in terms\n",
    "of time saved in your own projects"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "7Leri_zBe0Yy"
   },
   "source": [
    "# YOUR ACTION REQUIRED:\n",
    "# 1. Fork the MNIST example.\n",
    "# 2. Launch on Xmanager.\n",
    "# 3. Check out the Colab.\n",
    "# 4. Replace the dataset with the fashion mnist dataset from above.\n",
    "# 5. Re-run all tests to and fix if necessary.\n",
    "# 6. Launch modified version on Xmanager.\n",
    "# 6. Run Colab again with your updated code."
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "KJCf7xbhfWuH"
   },
   "source": [
    "# YOUR ACTION REQUIRED:\n",
    "# Check out the code of the MNIST and imagenet examples.\n",
    "# What differences do you see?"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ug1R0cMJfNox"
   },
   "source": [
    "### Mini project ?\n",
    "\n",
    "You might want to use brain templates and Flax examples for your AIR mini\n",
    "project.\n",
    "\n",
    "Suggestions\n",
    "\n",
    "- The go/bt-imagenet incorporates many best practices and might be a good\n",
    "  candidate to start your project from.\n",
    "- You might also want to have a look at the current Flax examples at\n",
    "  https://github.com/google/flax/tree/master/linen_examples\n",
    "  as a starting point for your project. You could try extracting the model code\n",
    "  from an example there and merging it into the brain template.\n",
    "\n",
    "If you encounter any problems on the way, you can reach us via go/flaxers-chat."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "oBgGdfd9O-2J"
   },
   "source": [
    "### end"
   ]
  }
 ]
}