<commit_before>{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "j_LlXHYcmRaC"
      },
      "source": [
        "# Gradient Accumulation"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vQro77whXULU"
      },
      "source": [
        "_Gradient accumulation_ is a technique where the gradients for several consecutive optimization steps are combined together, so that they can be applied at regular repeating intervals.\n",
        "\n",
        "One example where this is useful is to simulate training with a larger batch size than would fit into the available device memory. Another example is in the context of multi-task learning, where batches for different tasks may be visited in a round-robin fashion. Gradient accumulation makes it possible to simulate training on one large batch containing all of the tasks together.\n",
        "\n",
        "In this example, we give an example of implementing gradient accumulation using `optax.MultiSteps`. We start by bringing in some imports and defining some type annotations."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9cu0kFNrnJj7"
      },
      "outputs": [],
      "source": [
        "import functools\n",
        "from typing import Callable, Iterable, TypedDict\n",
        "\n",
        "import haiku as hk\n",
        "import jax\n",
        "import jax.numpy as jnp\n",
        "import numpy as np\n",
        "import optax\n",
        "import chex\n",
        "\n",
        "\n",
        "class MiniBatch(TypedDict):\n",
        "  image: jnp.ndarray\n",
        "  label: jnp.ndarray\n",
        "\n",
        "\n",
        "UpdateFn = Callable[[hk.Params, optax.OptState, MiniBatch],\n",
        "                    tuple[hk.Params, optax.OptState]]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TIKfATeXnW3B"
      },
      "source": [<commit_msg>Ensure colab example supports older version of python.<commit_after>{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "j_LlXHYcmRaC"
      },
      "source": [
        "# Gradient Accumulation"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vQro77whXULU"
      },
      "source": [
        "_Gradient accumulation_ is a technique where the gradients for several consecutive optimization steps are combined together, so that they can be applied at regular repeating intervals.\n",
        "\n",
        "One example where this is useful is to simulate training with a larger batch size than would fit into the available device memory. Another example is in the context of multi-task learning, where batches for different tasks may be visited in a round-robin fashion. Gradient accumulation makes it possible to simulate training on one large batch containing all of the tasks together.\n",
        "\n",
        "In this example, we give an example of implementing gradient accumulation using `optax.MultiSteps`. We start by bringing in some imports and defining some type annotations."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9cu0kFNrnJj7"
      },
      "outputs": [],
      "source": [
        "import functools\n",
        "from typing import Callable, Iterable, Tuple, TypedDict\n",
        "\n",
        "import haiku as hk\n",
        "import jax\n",
        "import jax.numpy as jnp\n",
        "import numpy as np\n",
        "import optax\n",
        "import chex\n",
        "\n",
        "\n",
        "class MiniBatch(TypedDict):\n",
        "  image: jnp.ndarray\n",
        "  label: jnp.ndarray\n",
        "\n",
        "\n",
        "UpdateFn = Callable[[hk.Params, optax.OptState, MiniBatch],\n",
        "                    Tuple[hk.Params, optax.OptState]]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TIKfATeXnW3B"
      },
      "source": [
{'fraction_non_alphanumeric': 0.12106299212598425, 'fraction_numerical': 0.002952755905511811, 'mean_word_length': 2.0795454545454546, 'pattern_counts': {'":': 38, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, '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': '2247399', 'n_tokens_mistral': 1234, 'n_tokens_neox': 1192, 'n_words': 386}