diff --git a/docs/notebooks/flax_guided_tour.ipynb b/docs/notebooks/flax_guided_tour.ipynb
index 567df9238..34b58cb21 100644
--- a/docs/notebooks/flax_guided_tour.ipynb
+++ b/docs/notebooks/flax_guided_tour.ipynb
@@ -1,41 +1,11 @@
 {
- "nbformat": 4,
- "nbformat_minor": 0,
- "metadata": {
-  "colab": {
-   "name": "A guided tour of Flax.ipynb",
-   "provenance": [],
-   "collapsed_sections": [],
-   "last_runtime": {
-    "build_target": "//learning/deepmind/dm_python:dm_notebook3_tpu",
-    "kind": "private"
-   },
-   "toc_visible": true
-  },
-  "kernelspec": {
-   "name": "python3",
-   "language": "python",
-   "display_name": "Python 3"
-  },
-  "pycharm": {
-   "stem_cell": {
-    "cell_type": "raw",
-    "source": [],
-    "metadata": {
-     "collapsed": false
-    }
-   }
-  }
- },
  "cells": [
   {
    "cell_type": "markdown",
-   "execution_count": null,
    "metadata": {
     "colab_type": "text",
     "id": "1ubnMn66ga4Z"
    },
-   "outputs": [],
    "source": [
     "# A guided tour of Flax\n",
     "\n",
@@ -64,12 +34,10 @@
   },
   {
    "cell_type": "markdown",
-   "execution_count": null,
    "metadata": {
     "colab_type": "text",
     "id": "HMNMi98ejnll"
    },
-   "outputs": [],
    "source": [
     "## Intro to JAX\n",
     "\n",
@@ -117,12 +85,10 @@
   },
   {
    "cell_type": "markdown",
-   "execution_count": null,
    "metadata": {
     "colab_type": "text",
     "id": "hwwff8vUJ6t7"
    },
-   "outputs": [],
    "source": [
     "Functional programming without abstractions naturally results into somewhat verbose but very explicit code.\n",
     "\n",
@@ -143,12 +109,10 @@
   },
   {
    "cell_type": "markdown",
-   "execution_count": null,
    "metadata": {
     "colab_type": "text",
     "id": "uiEnIech21gc"
    },
-   "outputs": [],
    "source": [
     "Once we generated a set of parameters it is easy enough to apply them to some inputs."
    ]
@@ -165,12 +129,10 @@
   },
   {
    "cell_type": "markdown",
-   "execution_count": null,
    "metadata": {
     "colab_type": "text",
     "id": "VPKvLaZl3BP1"
    },
-   "outputs": [],
    "source": [
     "Because everything is functional we can use the functional transformations that JAX provides to do useful things like taking gradients to optimize the model."
    ]
@@ -190,12 +152,10 @@
   },
   {
    "cell_type": "markdown",
-   "execution_count": null,
    "metadata": {
     "colab_type": "text",
     "id": "__CrzjEfqE5L"
    },
-   "outputs": [],
    "source": [
     "## Simplifying Neural Networks in JAX: Flax Modules\n",
     "\n",
@@ -262,12 +222,10 @@
   },
   {
    "cell_type": "markdown",
-   "execution_count": null,
    "metadata": {
     "colab_type": "text",
     "id": "ednxR882N4kX"
    },
-   "outputs": [],
    "source": [
     "Note that both `init` and `call` end up using the same `apply` function. That is why we must specify all the inputs\n",
     "and parameters (the number of features) in both `init` and `call`. Often the parameters are the same for each\n",
@@ -290,12 +248,10 @@
   },
   {
    "cell_type": "markdown",
-   "execution_count": null,
    "metadata": {
     "colab_type": "text",
     "id": "3Q1oe0WhOkxC"
    },
-   "outputs": [],
    "source": [
     "### Composition\n",
     "\n",
@@ -334,12 +290,10 @@
   },
   {
    "cell_type": "markdown",
-   "execution_count": null,
    "metadata": {
     "colab_type": "text",
     "id": "o6CPJofni4vf"
    },
-   "outputs": [],
    "source": [
     "The `params` returned by `init` have a nested structure of lists, tuples, dicts and \n",
     "and other types that can contain arrays; we call such a structure a pytree.\n",
@@ -359,12 +313,10 @@
   },
   {
    "cell_type": "markdown",
-   "execution_count": null,
    "metadata": {
     "colab_type": "text",
     "id": "OGyMhq1rmBRr"
    },
-   "outputs": [],
    "source": [
     "#### Module name\n",
     "By default Flax will use integers as keys for the parameters of sub Modules. By passing the `name` argument we can control the parameter structure and make it more meaningful."
@@ -394,12 +346,10 @@
   },
   {
    "cell_type": "markdown",
-   "execution_count": null,
    "metadata": {
     "colab_type": "text",
     "id": "_L5AOQ97mZAA"
    },
-   "outputs": [],
    "source": [
     "### Parameter sharing\n",
     "\n",
@@ -429,12 +379,10 @@
   },
   {
    "cell_type": "markdown",
-   "execution_count": null,
    "metadata": {
     "colab_type": "text",
     "id": "Bn396SL6RARD"
    },
-   "outputs": [],
    "source": [
     "we call the Dense layer named 'cell' 3 times but only one set of parameters shows up in the parameter structure due to weight sharing."
    ]
@@ -452,12 +400,10 @@
   },
   {
    "cell_type": "markdown",
-   "execution_count": null,
    "metadata": {
     "colab_type": "text",
     "id": "lFUvB2IcZJpu"
    },
-   "outputs": [],
    "source": [
     "### Shape inference\n",
     "\n",
@@ -484,12 +430,10 @@
   },
   {
    "cell_type": "markdown",
-   "execution_count": null,
    "metadata": {
     "colab_type": "text",
     "id": "EAdFRvBwRran"
    },
-   "outputs": [],
    "source": [
     "### Model\n",
     "\n",
@@ -534,12 +478,10 @@
   },
   {
    "cell_type": "markdown",
-   "execution_count": null,
    "metadata": {
     "colab_type": "text",
     "id": "CDvnJ1VP_qfz"
    },
-   "outputs": [],
    "source": [
     "Parameters can be updated using the `Model.replace` method"
    ]
@@ -556,12 +498,10 @@
   },
   {
    "cell_type": "markdown",
-   "execution_count": null,
    "metadata": {
     "colab_type": "text",
     "id": "G8PAwSzbA29M"
    },
-   "outputs": [],
    "source": [
     "Model is registered as a JAX pytree container object which means that it can be passed to JAX transformations and `jax.tree_map`.\n",
     "\n",
@@ -584,12 +524,10 @@
   },
   {
    "cell_type": "markdown",
-   "execution_count": null,
    "metadata": {
     "colab_type": "text",
     "id": "yHxj5gVhAzLX"
    },
-   "outputs": [],
    "source": [
     "### State\n",
     "\n",
@@ -645,12 +583,10 @@
   },
   {
    "cell_type": "markdown",
-   "execution_count": null,
    "metadata": {
     "colab_type": "text",
     "id": "QKoNxlq84lHk"
    },
-   "outputs": [],
    "source": [
     "Stateful modules require special care when used. The `nn.stateful` context manager defines a scope in which stateful operations are allowed. Outside of this scope the state becomes immutable.\n",
     "\n",
@@ -694,12 +630,10 @@
   },
   {
    "cell_type": "markdown",
-   "execution_count": null,
    "metadata": {
     "colab_type": "text",
     "id": "M2DJaEVV6I67"
    },
-   "outputs": [],
    "source": [
     "The state can be inspected using `Collection.as_dict()`.\n",
     "\n",
@@ -726,12 +660,10 @@
   },
   {
    "cell_type": "markdown",
-   "execution_count": null,
    "metadata": {
     "colab_type": "text",
     "id": "NWqk1SBZ6iWD"
    },
-   "outputs": [],
    "source": [
     "The stateful mechanism forces the user to be explicit about stateful operations.\n",
     "\n",
@@ -758,12 +690,10 @@
   },
   {
    "cell_type": "markdown",
-   "execution_count": null,
    "metadata": {
     "colab_type": "text",
     "id": "nWdH5jcr7d05"
    },
-   "outputs": [],
    "source": [
     "We can solve the skew in statistics by creating a separate state for this alternative input distribution."
    ]
@@ -785,12 +715,10 @@
   },
   {
    "cell_type": "markdown",
-   "execution_count": null,
    "metadata": {
     "colab_type": "text",
     "id": "KEfIkD3uR8hh"
    },
-   "outputs": [],
    "source": [
     "## Optimizer\n",
     "\n",
@@ -819,9 +747,7 @@
   },
   {
    "cell_type": "markdown",
-   "execution_count": null,
    "metadata": {},
-   "outputs": [],
    "source": [
     "The model is nothing more than a Dense module with a single feature"
    ]
@@ -850,9 +776,7 @@
   },
   {
    "cell_type": "markdown",
-   "execution_count": null,
    "metadata": {},
-   "outputs": [],
    "source": [
     "We will use gradient descent with momentum to fit the model to the data.\n",
     "Each optimizer inherits from the `flax.optim.OptimizerDef`.\n",
@@ -881,7 +805,7 @@
     "\n",
     "for i in range(train_steps):\n",
     "  # optimizer.target is passed to the loss_fn\n",
-    "  loss, grad = optimizer.compute_gradient(loss_fn)\n",
+    "  loss, grad = jax.value_and_grad(loss_fn)(optimizer.target)\n",
     "  # `apply_gradient` returns a new `Optimizer` instance with the updated target and optimizer state.\n",
     "  optimizer = optimizer.apply_gradient(grad)\n",
     "print('mean square error:', loss)\n",
@@ -896,9 +820,7 @@
   },
   {
    "cell_type": "markdown",
-   "execution_count": null,
    "metadata": {},
-   "outputs": [],
    "source": [
     "## Serialization\n",
     "\n",
@@ -926,9 +848,7 @@
   },
   {
    "cell_type": "markdown",
-   "execution_count": null,
    "metadata": {},
-   "outputs": [],
    "source": [
     "Flax dataclasses are immutable. Using the `replace` method a new instance can be created with a set of updated fields."
    ]
@@ -945,9 +865,7 @@
   },
   {
    "cell_type": "markdown",
-   "execution_count": null,
    "metadata": {},
-   "outputs": [],
    "source": [
     "The `to_bytes` and `from_bytes` functions are used to convert an object to the message pack format and back."
    ]
@@ -985,24 +903,20 @@
   },
   {
    "cell_type": "markdown",
-   "execution_count": null,
    "metadata": {
     "colab_type": "text",
     "id": "1KCPh73QSsLr"
    },
-   "outputs": [],
    "source": [
     "## Advanced features"
    ]
   },
   {
    "cell_type": "markdown",
-   "execution_count": null,
    "metadata": {
     "colab_type": "text",
     "id": "jJyqXS1RSvXg"
    },
-   "outputs": [],
    "source": [
     "### Selective Optimization\n",
     "\n",
@@ -1038,7 +952,7 @@
     "  return jnp.square(Y - Y_hat).mean()\n",
     "\n",
     "for i in range(train_steps):\n",
-    "  loss, grad = optimizer.compute_gradient(loss_fn)\n",
+    "  loss, grad = jax.value_and_grad(loss_fn)(optimizer.target)\n",
     "  optimizer = optimizer.apply_gradient(grad)\n",
     "print('mean square error:', loss)\n",
     "\n",
@@ -1052,12 +966,10 @@
   },
   {
    "cell_type": "markdown",
-   "execution_count": null,
    "metadata": {
     "colab_type": "text",
     "id": "CHQhNahuS0ed"
    },
-   "outputs": [],
    "source": [
     "### Multi method modules\n",
     "\n",
@@ -1106,12 +1018,10 @@
   },
   {
    "cell_type": "markdown",
-   "execution_count": null,
    "metadata": {
     "colab_type": "text",
     "id": "zbA5RyHVTC_V"
    },
-   "outputs": [],
    "source": [
     "### Transforming sub module parameters"
    ]
@@ -1157,6 +1067,45 @@
    "outputs": [],
    "source": []
   }
- ]
+ ],
+ "metadata": {
+  "colab": {
+   "collapsed_sections": [],
+   "last_runtime": {
+    "build_target": "//learning/deepmind/dm_python:dm_notebook3_tpu",
+    "kind": "private"
+   },
+   "name": "A guided tour of Flax.ipynb",
+   "provenance": [],
+   "toc_visible": true
+  },
+  "kernelspec": {
+   "display_name": "Flax",
+   "language": "python",
+   "name": "flax"
+  },
+  "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.6.10"
+  },
+  "pycharm": {
+   "stem_cell": {
+    "cell_type": "raw",
+    "metadata": {
+     "collapsed": false
+    },
+    "source": []
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
 }
-