synopsis.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4CldxEhqQac_"
   },
   "source": [
    "# JAX Loop Utils\n",
    "\n",
    "<a href=\"https://colab.research.google.com/github/Astera-org/jax_loop_utils/blob/master/synopsis.ipynb\" ><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>\n",
    "\n",
    "\n",
    "    pip install git+https://github.com/Astera-org/jax_loop_utils\n",
    "\n",
    "https://github.com/Astera-org/jax_loop_utils\n",
    "\n",
    "Writing and maintaining your own training loop gives lots of\n",
    "flexibility but also quickly leads to non-trivial amount of code that is\n",
    "repeated in every project.\n",
    "\n",
    "`jax_loop_utils` provides small independent helpers to make the training loop shorter and\n",
    "easier to read, while keeping maximum flexibility.\n",
    "\n",
    "**This Colab** walks you through the different modules of `jax_loop_utils` with simple\n",
    "example code for showcasing the important concepts and to be pasted into your\n",
    "training loop to get started using `jax_loop_utils`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6CxpJUPTd1wD"
   },
   "source": [
    "### Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "sTS7uCbgizD8"
   },
   "outputs": [],
   "source": [
    "%pip install -q \"git+https://github.com/Astera-org/jax_loop_utils[tf-data]\" flax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "6wmDlgxtdsQ7",
    "outputId": "6cdd353c-165e-4b88-efc4-8dd46b01c331"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Name: jax-loop-utils\n",
      "Version: 0.0.12\n",
      "Location: /Users/garymm/src/Astera-org/jax_loop_utils/.venv/lib/python3.12/site-packages\n",
      "Editable project location: /Users/garymm/src/Astera-org/jax_loop_utils\n",
      "Requires: absl-py, etils, jax, jaxlib, ml-collections, numpy, packaging, typing-extensions, wrapt\n",
      "Required-by:\n"
     ]
    }
   ],
   "source": [
    "%pip show jax-loop-utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "WrGls5mON3Qr"
   },
   "outputs": [],
   "source": [
    "import jax\n",
    "import jax.numpy as jnp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "-kx9lHndyN8G"
   },
   "outputs": [],
   "source": [
    "import chex\n",
    "\n",
    "chex.set_n_cpu_devices(2)  # Simulate 2 local devices in a CPU Colab runtime."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "C83tJrZnzTns"
   },
   "source": [
    "### `jax_loop_utils.metric_writers`\n",
    "\n",
    "The module [`metric_writers`] provides a simple [interface] to write time series\n",
    "metrics in a unified way.\n",
    "\n",
    "Metric writers provided:\n",
    "\n",
    "- `SummaryWriter`: Uses `tf.summary` to write summary files. For display in\n",
    "  TensorBoard.\n",
    "- `LoggingWriter`: Simply writes values to the INFO log. This obviously only\n",
    "  supports data types that can be converted to text but is still helpful for\n",
    "  seeing the training progress on the command line.\n",
    "- `TorchTensorboardWriter`: Uses `torch.utils.tensorboard` to write summary\n",
    "  files. Use this writer for the Pytorch-based code.\n",
    "\n",
    "Additional we provide metric writers to combine multiple metric writers\n",
    "(`MultiWriter`) and to move the write operation to a background thread\n",
    "(`AsyncWriter`).\n",
    "\n",
    "[`metric_writers`]: https://github.com/Astera-org/jax_loop_utils/blob/master/jax_loop_utils/metric_writers/__init__.py\n",
    "[interface]: https://github.com/Astera-org/jax_loop_utils/blob/master/jax_loop_utils/metric_writers/interface.py\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "id": "aBhPGVLKz8c6"
   },
   "outputs": [],
   "source": [
    "from absl import logging\n",
    "\n",
    "logging.set_verbosity(logging.INFO)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "id": "uO_ezNP216XM"
   },
   "outputs": [],
   "source": [
    "logdir = \"./metrics\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Jt1eLAfNz_wr",
    "outputId": "70668139-45cc-416a-ff3d-d8d99aece0a8"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:absl:[0] loss=1\n",
      "INFO:absl:[1] loss=0.9\n",
      "INFO:absl:[2] loss=0.81\n",
      "INFO:absl:[3] loss=0.729\n",
      "INFO:absl:[4] loss=0.6561\n",
      "INFO:absl:[5] loss=0.59049\n",
      "INFO:absl:[6] loss=0.531441\n",
      "INFO:absl:[7] loss=0.478297\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:absl:[8] loss=0.430467\n",
      "INFO:absl:[9] loss=0.38742\n",
      "INFO:absl:[10] steps_per_sec=459770\n",
      "INFO:absl:[10] uptime=0.0005575\n",
      "INFO:absl:[20] steps_per_sec=439.342\n",
      "INFO:absl:[20] uptime=0.0238017\n",
      "INFO:absl:[30] steps_per_sec=9074.06\n",
      "INFO:absl:[30] uptime=0.0257181\n",
      "INFO:absl:[40] steps_per_sec=5392.65\n",
      "INFO:absl:[40] uptime=0.028389\n",
      "INFO:absl:[50] steps_per_sec=3739.37\n",
      "INFO:absl:[50] uptime=0.0304817\n",
      "INFO:absl:[60] steps_per_sec=4778.5\n",
      "INFO:absl:[60] uptime=0.0324869\n",
      "INFO:absl:[70] steps_per_sec=5000.73\n",
      "INFO:absl:[70] uptime=0.0336201\n",
      "INFO:absl:[80] steps_per_sec=8857.07\n",
      "INFO:absl:[80] uptime=0.0361314\n",
      "INFO:absl:[90] steps_per_sec=3984.86\n",
      "INFO:absl:[90] uptime=0.0388259\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import pathlib\n",
    "from jax_loop_utils import metric_writers\n",
    "from jax_loop_utils.metric_writers.tf import SummaryWriter\n",
    "\n",
    "# Handy shortcut to create create async logging/tensorboard writer.\n",
    "\n",
    "\n",
    "def create_default_writer(\n",
    "    logdir: os.PathLike | None = None,\n",
    "    *,\n",
    "    just_logging: bool = False,\n",
    "    asynchronous: bool = True,\n",
    "    collection: str | None = None,\n",
    ") -> metric_writers.MultiWriter:\n",
    "    \"\"\"Create the default writer for the platform.\n",
    "\n",
    "    On most platforms this will create a MultiWriter that writes to multiple back\n",
    "    ends (logging, TF summaries etc.).\n",
    "\n",
    "    Args:\n",
    "      logdir: Logging dir to use for TF summary files. If empty/None will the\n",
    "        returned writer will not write TF summary files.\n",
    "      just_logging: If True only use a LoggingWriter. This is useful in multi-host\n",
    "        setups when only the first host should write metrics and all other hosts\n",
    "        should only write to their own logs.\n",
    "        default (None) will automatically determine if you # GOOGLE-INTERNAL have\n",
    "      asynchronous: If True return an AsyncMultiWriter to not block when writing\n",
    "        metrics.\n",
    "      collection: A string which, if provided, provides an indication that the\n",
    "        provided metrics should all be written to the same collection, or\n",
    "        grouping.\n",
    "\n",
    "    Returns:\n",
    "      A `MetricWriter` according to the platform and arguments.\n",
    "    \"\"\"\n",
    "    if just_logging:\n",
    "        if asynchronous:\n",
    "            return metric_writers.AsyncMultiWriter(\n",
    "                [metric_writers.LoggingWriter(collection=collection)]\n",
    "            )\n",
    "        else:\n",
    "            return metric_writers.MultiWriter(\n",
    "                [metric_writers.LoggingWriter(collection=collection)]\n",
    "            )\n",
    "    writers = [metric_writers.LoggingWriter(collection=collection)]\n",
    "    if logdir is not None:\n",
    "        logdir = pathlib.Path(logdir)\n",
    "        if collection is not None:\n",
    "            logdir /= collection\n",
    "        writers.append(SummaryWriter(os.fspath(logdir)))\n",
    "    if asynchronous:\n",
    "        return metric_writers.AsyncMultiWriter(writers)\n",
    "    return metric_writers.MultiWriter(writers)\n",
    "\n",
    "\n",
    "writer = create_default_writer(logdir)\n",
    "for step in range(10):\n",
    "    writer.write_scalars(step, dict(loss=0.9**step))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "9hEe8xAO0M-1"
   },
   "outputs": [],
   "source": [
    "%load_ext tensorboard\n",
    "%tensorboard --logdir=./metrics"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "VyTm-xlM0Zvz"
   },
   "source": [
    "### `jax_loop_utils.periodic_actions`\n",
    "\n",
    "[`periodic_actions`] are simple helpers that allow you to do in the training\n",
    "loop at regular intervals. Currently we support\n",
    "\n",
    "- `PeriodicAction`, `PeriodicCallback`: To implement your own actions.\n",
    "- `Profile`: To create TensorBoard compatible profiles.\n",
    "- `ReportProgress`: To continuously print progress status updates.\n",
    "\n",
    "[`periodic_actions`]: https://github.com/Astera-org/jax_loop_utils/blob/master/jax_loop_utils/periodic_actions.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "vsKd1Frm1cRN",
    "outputId": "da99a807-b861-464a-f13c-8fe76049b749"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:absl:Setting work unit notes: 459770.3 steps/s, 10.0% (10/100), ETA: 0m\n",
      "INFO:absl:Setting work unit notes: 439.3 steps/s, 20.0% (20/100), ETA: 0m\n",
      "INFO:absl:Setting work unit notes: 9074.1 steps/s, 30.0% (30/100), ETA: 0m\n",
      "INFO:absl:Setting work unit notes: 5392.7 steps/s, 40.0% (40/100), ETA: 0m\n",
      "INFO:absl:Setting work unit notes: 3739.4 steps/s, 50.0% (50/100), ETA: 0m\n",
      "INFO:absl:Setting work unit notes: 4778.5 steps/s, 60.0% (60/100), ETA: 0m\n",
      "INFO:absl:Setting work unit notes: 5000.7 steps/s, 70.0% (70/100), ETA: 0m\n",
      "INFO:absl:Setting work unit notes: 8857.1 steps/s, 80.0% (80/100), ETA: 0m\n",
      "INFO:absl:Setting work unit notes: 3984.9 steps/s, 90.0% (90/100), ETA: 0m\n"
     ]
    }
   ],
   "source": [
    "from jax_loop_utils import periodic_actions\n",
    "\n",
    "total_steps = 100\n",
    "hooks = [\n",
    "    # Outputs progress via metric writer (in this case logs & TensorBoard).\n",
    "    periodic_actions.ReportProgress(\n",
    "        num_train_steps=total_steps, every_steps=10, writer=writer\n",
    "    ),\n",
    "    periodic_actions.Profile(logdir=logdir),\n",
    "]\n",
    "\n",
    "for step in range(total_steps):\n",
    "    for hook in hooks:\n",
    "        hook(step)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "uB7-MTls1-No",
    "outputId": "bf1bd0b3-48e0-475c-81df-6dc9ebb0ce6a"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "total 8\n",
      "-rw-r--r--  1 garymm  staff   1.9K Nov 15 16:41 events.out.tfevents.1731717685.Garys-MacBook-Pro.local.3562.0.v2\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/homebrew/Cellar/python@3.12/3.12.4/Frameworks/Python.framework/Versions/3.12/lib/python3.12/pty.py:95: RuntimeWarning: os.fork() was called. os.fork() is incompatible with multithreaded code, and JAX is multithreaded, so this will likely lead to a deadlock.\n",
      "  pid, fd = os.forkpty()\n"
     ]
    }
   ],
   "source": [
    "# If you click on \"refresh\" in above TensorBoard you'll now see a new\n",
    "# \"steps_per_sec\" metric...\n",
    "!ls -lh metrics"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ueom-uBWLbeQ"
   },
   "source": [
    "### `jax_loop_utils.metrics`\n",
    "\n",
    "The [`metrics`] module provides a framework for functional metric computation.\n",
    "Note that this module does **not** include the actual metric definitions (other\n",
    "than `metrics.Accuracy` that is provided for demonstration purposes), but\n",
    "rather provides abstractions that can be used to compute metrics in a\n",
    "distributed distributed environment.\n",
    "\n",
    "This section is a bit longer than the previous sections and walks you through\n",
    "the following parts:\n",
    "\n",
    "1. How `metrics.Metric` is computed, and defining \"averageable\" metrics.\n",
    "2. Using `metrics.Collection` to compute several metrics at once.\n",
    "3. Aggregating in an evaluation step that is transformed by `pmap()`.\n",
    "4. Define a new metric with custom aggregation (i.e. non \"averageable\").\n",
    "\n",
    "\n",
    "[`metrics`]: https://github.com/Astera-org/jax_loop_utils/blob/master/jax_loop_utils/metrics.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Bti6-K5ZDZHP",
    "outputId": "701eb451-8d32-4b11-e389-4599c41c581c"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(0.75, dtype=float32)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from jax_loop_utils import metrics\n",
    "import flax\n",
    "\n",
    "# Metrics are computed in three steps:\n",
    "\n",
    "# 1. Compute intermediate values from model outputs\n",
    "accuracy_batch1 = metrics.Accuracy.from_model_output(\n",
    "    logits=jnp.array([[-1.0, 1.0], [1.0, -1.0]]),\n",
    "    labels=jnp.array([0, 0]),  # i.e. 1st incorrect, 2nd correct\n",
    ")\n",
    "accuracy_batch2 = metrics.Accuracy.from_model_output(\n",
    "    logits=jnp.array([[-1.0, 1.0], [1.0, -1.0]]),\n",
    "    labels=jnp.array([1, 0]),  # i.e. both correct\n",
    ")\n",
    "\n",
    "# 2. Intermediate values are aggregated\n",
    "accuracy = accuracy_batch1\n",
    "accuracy = accuracy.merge(accuracy_batch2)\n",
    "\n",
    "# 3. Final metrics are computed from aggregated intermediate values:\n",
    "accuracy.compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "ga8J-z6KmMBg",
    "outputId": "42dd6d88-d40e-46ea-b4f0-cba649784924"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(2.2, dtype=float32)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# It's easy to define your own metrics if they are \"averageable\":\n",
    "\n",
    "AverageLoss = metrics.Average.from_output(\"loss\")\n",
    "\n",
    "AverageLoss.from_model_output(loss=jnp.array([1.1, 3.3])).compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "PsxIu5FfnJmW",
    "outputId": "a5ecc19c-6772-4ca6-f442-52cde0bc1f57"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(2.1999998, dtype=float32)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# You can provide a functional to derive the value-to-be-averaged:\n",
    "\n",
    "# Note that our metric only uses the model output named \"loss\". There can be an\n",
    "# arbitrary number of additional model outputs that we don't need here (**_).\n",
    "AverageSquaredLoss = metrics.Average.from_fun(lambda loss, **_: loss**2)\n",
    "\n",
    "AverageSquaredLoss.from_model_output(loss=jnp.array([1.1**0.5, 3.3**0.5])).compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "zuXskYw_ooKQ",
    "outputId": "916be6e3-916a-4070-d00f-388ca29a16f7"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'loss': Array(2.2, dtype=float32), 'accuracy': Array(0.75, dtype=float32)}"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Usually you would want to compute a collection of metrics from model outputs:\n",
    "\n",
    "\n",
    "@flax.struct.dataclass  # <-- required for JAX transformations\n",
    "class MyMetrics(metrics.Collection):\n",
    "    loss: metrics.Average.from_output(\"loss\")  # type: ignore[invalid-type-form]\n",
    "    accuracy: metrics.Accuracy\n",
    "\n",
    "\n",
    "# 1. Compute intermediate values from model outputs\n",
    "my_metrics_batch1 = MyMetrics.single_from_model_output(\n",
    "    logits=jnp.array([[-1.0, 1.0], [1.0, -1.0]]),\n",
    "    labels=jnp.array([0, 0]),  # i.e. 1st incorrect, 2nd correct\n",
    "    loss=jnp.array([3.3, 2.2]),\n",
    ")\n",
    "my_metrics_batch2 = MyMetrics.single_from_model_output(\n",
    "    logits=jnp.array([[-1.0, 1.0], [1.0, -1.0]]),\n",
    "    labels=jnp.array([1, 0]),  # i.e. both correct\n",
    "    loss=jnp.array([2.2, 1.1]),\n",
    ")\n",
    "\n",
    "# 2. Intermediate values are aggregated\n",
    "my_metrics = my_metrics_batch1.merge(my_metrics_batch2)\n",
    "\n",
    "# 3. Final metrics are computed from aggregated intermediate values:\n",
    "my_metrics.compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "8v7vxEGGqArt",
    "outputId": "b4d77b53-1032-4db8-c275-65aefb558c67"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'loss': Array(2.2, dtype=float32), 'accuracy': Array(0.75, dtype=float32)}"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Often you want to compute these metrics inside a pmap(). The framework\n",
    "# provides the handy `Collection.gather_from_model_output` that will first\n",
    "# compute the intermediate values, then call `jax.lax.all_gather()` to gather\n",
    "# the intermediate values from all the devices (in a multi-host setup that's\n",
    "# all the devices in the mesh, not only the local devices), and then reduce them\n",
    "# by calling `Metric.merge()` in a `jax.lax.scan()` loop.\n",
    "\n",
    "# Sounds complicated? Using it is actually surprisingly simple:\n",
    "\n",
    "\n",
    "def fake_model(params, batch):\n",
    "    del params  # Fake.\n",
    "    return batch\n",
    "\n",
    "\n",
    "def eval_step(my_metrics, params, batch):\n",
    "    model_outputs = fake_model(params, batch)\n",
    "    # IMPORTANT: If you called `.single_from_model_output()` here, then all values\n",
    "    # from devices after the first device would be ignored for the metric\n",
    "    # computation.\n",
    "    return my_metrics.merge(MyMetrics.gather_from_model_output(**model_outputs))\n",
    "\n",
    "\n",
    "eval_step_p = jax.pmap(eval_step, axis_name=\"batch\")\n",
    "\n",
    "my_metrics = flax.jax_utils.replicate(MyMetrics.empty())\n",
    "\n",
    "for batch in [\n",
    "    # Single batch of data pmapped on two devices in parallel.\n",
    "    dict(\n",
    "        logits=jnp.array(\n",
    "            [\n",
    "                # Batch for device 1\n",
    "                [[-1.0, 1.0], [1.0, -1.0]],\n",
    "                # Batch for device 2\n",
    "                [[-1.0, 1.0], [1.0, -1.0]],\n",
    "            ]\n",
    "        ),\n",
    "        labels=jnp.array(\n",
    "            [\n",
    "                # Batch for device 1\n",
    "                [0, 0],\n",
    "                # Batch for device 2\n",
    "                [1, 0],\n",
    "            ]\n",
    "        ),\n",
    "        loss=jnp.array(\n",
    "            [\n",
    "                # Batch for device 1\n",
    "                [3.3, 2.2],\n",
    "                # Batch for device 2\n",
    "                [2.2, 1.1],\n",
    "            ]\n",
    "        ),\n",
    "    ),\n",
    "]:\n",
    "    my_metrics = eval_step_p(my_metrics, None, batch)\n",
    "\n",
    "# Note that up to this point all inputs/outputs to `eval_step_p()` are\n",
    "# replicated such that their leading dimension == number of local devices == 8.\n",
    "my_metrics.unreplicate().compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "fPKbjXgVOxnN",
    "outputId": "edc422ce-d8d2-4450-e7fe-439e1b0dbe9c"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Note that not calling `.unreplicate()` raises an erorr: Collection is still replicated (ndim=1). Maybe you forgot to call a flax.jax_utils.unreplicate() or a Collections.reduce()?\n"
     ]
    }
   ],
   "source": [
    "try:\n",
    "    my_metrics.compute()\n",
    "    raise RuntimeError(\"Expected ValueError!\")\n",
    "except ValueError as e:\n",
    "    print(\"Note that not calling `.unreplicate()` raises an erorr:\", e)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "8v0gvSp3m9uW",
    "outputId": "58e3e3a9-7e62-422e-b09a-f93023b2edec"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(0.6666667, dtype=float32)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# You can also provide your own aggregation logic:\n",
    "\n",
    "\n",
    "@flax.struct.dataclass\n",
    "class Precision(metrics.Metric):\n",
    "    \"\"\"Computes the precision from model outputs `logits` and `labels`.\"\"\"\n",
    "\n",
    "    true_positives: jnp.array\n",
    "    pred_positives: jnp.array\n",
    "\n",
    "    @classmethod\n",
    "    def from_model_output(\n",
    "        cls, *, logits: jnp.array, labels: jnp.array, **_\n",
    "    ) -> metrics.Metric:\n",
    "        assert logits.shape[-1] == 2, \"Expected binary logits.\"\n",
    "        preds = logits.argmax(axis=-1)\n",
    "        return cls(\n",
    "            true_positives=((preds == 1) & (labels == 1)).sum(),\n",
    "            pred_positives=(preds == 1).sum(),\n",
    "        )\n",
    "\n",
    "    def merge(self, other: metrics.Metric) -> metrics.Metric:\n",
    "        # Note that for precision we cannot average metric values because the\n",
    "        # denominator of the metric value is pred_positives and not every batch of\n",
    "        # examples has the same number of pred_positives (as opposed to e.g.\n",
    "        # accuracy where every batch has the same number of)\n",
    "        return type(self)(\n",
    "            true_positives=self.true_positives + other.true_positives,\n",
    "            pred_positives=self.pred_positives + other.pred_positives,\n",
    "        )\n",
    "\n",
    "    def compute(self):\n",
    "        return self.true_positives / self.pred_positives\n",
    "\n",
    "\n",
    "Precision.from_model_output(\n",
    "    # 1 TP, 1 FN -- 2 pred_positives -- precision = 1.0\n",
    "    logits=jnp.array([[-1.0, 1.0], [1.0, -1.0]]),\n",
    "    labels=jnp.array([1, 1]),  # i.e. 1st incorrect, 2nd correct\n",
    ").merge(\n",
    "    Precision.from_model_output(\n",
    "        # 1 TP, 1 FP -- 2 pred_positives -- precision = 0.5\n",
    "        logits=jnp.array([[-1.0, 1.0], [-1.0, 1.0]]),\n",
    "        labels=jnp.array([1, 0]),  # i.e. 1st incorrect, 2nd correct\n",
    "    )\n",
    ").compute()\n",
    "\n",
    "# If one incorrectly used metrics.Average to aggregate the metric, the final\n",
    "# value would be 0.75 because both batches have the same weight in terms of\n",
    "# examples. But the first batch constains 2 pred_positives and should thus be\n",
    "# weighted 2x, resulting in the correct (1 + 1) / (1 + 2) == 0.66"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "name": "clu synopsis",
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "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.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}