diff --git a/metrax_example.ipynb b/metrax_example.ipynb new file mode 100644 index 0000000..59f17b1 --- /dev/null +++ b/metrax_example.ipynb @@ -0,0 +1,526 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": { + "id": "R_KB8l_ZKoa8" + }, + "source": [ + "Please connect to `Metrax (go/metrax)` colab runtime.\n", + "\n", + "If you dont see `Metrax (go/metrax)` from the dropdown menu, please run `/google/bin/releases/colaboratory/public/tools/authorize_colab` on your gLinux workstation or cloudtop and try again." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "LLecX4XsEoMH" + }, + "source": [ + "# Getting Started with Metrax šŸš€\n", + "\n", + "Welcome to this hands-on guide for `metrax`, a powerful and flexible metrics library for JAX.\n", + "\n", + "In this Colab, you'll learn how to:\n", + "* Use the **Functional metrax API** (`metrax`) and the **Object-Oriented metrax API** (`metrax.nnx`).\n", + "* Verify that batch and iterative calculations give **identical results**.\n", + "* Scale your metric computations to **multiple devices** using 1)`jax.pmap` and 2)`jax.jit`.\n", + "* Scale your metric computations to **multiple hosts** for large scale distributed training, using 1) `Multi-Controller JAX` and 2) `Pathways` solutions." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "mTwMjngoy-h8" + }, + "source": [ + "## āš™ļø Environment Setup: Simulating Multiple Devices\n", + "\n", + "First, let's configure our environment. To demonstrate `metrax`'s multi-device capabilities, we'll instruct JAX to simulate an environment with **4 virtual CPU devices**. This allows us to test `jax.pmap` and `jax.jit` with `mesh` logic even on single-device hardware." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "bhooimvEyzFc" + }, + "outputs": [], + "source": [ + "# This environment variable instructs JAX's underlying XLA compiler\n", + "# to create a specific number of virtual CPU devices.\n", + "#\n", + "# This MUST be set *before* the JAX backend is initialized, which happens on\n", + "# the first import of `jax`.\n", + "#\n", + "# In a script, ensure this line comes before `import jax`.\n", + "# In a notebook, a kernel restart may be needed if JAX has already been used.\n", + "import os\n", + "print(\"Configuring JAX to simulate 4 CPU devices...\")\n", + "os.environ['XLA_FLAGS'] = '--xla_force_host_platform_device_count=4'\n", + "import jax\n", + "\n", + "\n", + "# --- Verify the JAX Environment ---\n", + "print(\"\\nVerifying JAX environment configuration:\")\n", + "print(\"-\" * 40)\n", + "device_count = jax.device_count()\n", + "process_count = jax.process_count()\n", + "print(f\"āœ… Number of available JAX devices: {device_count}\")\n", + "print(f\"āœ… Number of JAX processes: {process_count}\")\n", + "print(\"-\" * 40)\n", + "\n", + "if device_count == 4:\n", + " print(\"Success! JAX is now set up for multi-device simulation.\")\n", + "else:\n", + " print(\"Warning: JAX device count is not the expected value.\")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "XjjUn5OVwhQW" + }, + "source": [ + "## šŸ“Š Data Preparation for Realistic Scenarios\n", + "\n", + "Next, let's generate some data. A good metrics demo uses realistic data, so we'll create a dataset that is **imbalanced** and where the model's **predictions are imperfect but correlated** with the true labels." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "ppl-mTUk6XWC" + }, + "outputs": [], + "source": [ + "import numpy as np\n", + "\n", + "# --- 1. Data Generation Setup ---\n", + "np.random.seed(42)\n", + "N_BATCHES = 4\n", + "BATCH_SIZE = 8\n", + "TOTAL_SAMPLES = N_BATCHES * BATCH_SIZE\n", + "\n", + "# --- 2. Create Realistic, Correlated Data ---\n", + "# Create an imbalanced dataset (80% class 0, 20% class 1).\n", + "labels = np.random.choice([0, 1], size=(TOTAL_SAMPLES,), p=[0.8, 0.2])\n", + "\n", + "# Generate predictions correlated with labels, adding some noise for realism.\n", + "noise = np.random.normal(loc=0, scale=0.25, size=TOTAL_SAMPLES)\n", + "clean_preds = np.where(labels == 1, 0.8, 0.2)\n", + "predictions = np.clip(clean_preds + noise, 0, 1)\n", + "\n", + "# Generate sample weights to give more importance to the rare positive class.\n", + "sample_weights = np.where(labels == 1, 2.0, 1.0)\n", + "\n", + "# --- 3. Reshape Data into Batched Format ---\n", + "# The batched format is useful for demonstrating iterative calculations.\n", + "labels_batched = labels.reshape(N_BATCHES, BATCH_SIZE).astype(np.float32)\n", + "predictions_batched = predictions.reshape(N_BATCHES, BATCH_SIZE).astype(np.float32)\n", + "sample_weights_batched = sample_weights.reshape(N_BATCHES, BATCH_SIZE).astype(np.float32)\n", + "\n", + "# --- 4. Data Shape Verification ---\n", + "print(\"āœ… Data generation complete. Verifying array shapes:\")\n", + "print(\"-\" * 50)\n", + "print(f\"Flat arrays for full-dataset processing:\")\n", + "print(f\" - predictions.shape: {predictions.shape}\")\n", + "print(f\" - labels.shape: {labels.shape}\")\n", + "print(f\"\\nBatched arrays for iterative/streaming processing:\")\n", + "print(f\" - predictions_batched.shape: {predictions_batched.shape}\")\n", + "print(f\" - labels_batched.shape: {labels_batched.shape}\")\n", + "print(\"-\" * 50)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "f4E5A--_cF86" + }, + "source": [ + "## Metrax Linen API (Functional)\n", + "\n", + "The core `metrax` API is functional and stateless, making it a natural fit for JAX. It works by creating immutable `Metric` state objects that can be merged.\n", + "\n", + "Each `metrax` metric inherits the CLU [`metric`](http://shortn/_e70RtO7j36) class and provides the following APIs:\n", + "\n", + "* `Metric.from_model_output()`: Creates a metric state from data.\n", + "* `Metric.empty()`: Creates an empty, initial state.\n", + "* `metric_a.merge(metric_b)`: Combines two metric states.\n", + "* `metric.compute()`: Computes the final value.\n", + "\n", + "Let's demonstrate by calculating several metrics on our dataset, once on the 1) full batch and once by 2) iteratively merging results. The second method resembles real world machine learning metrics calculation scenarios." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "eexpYXIELRUd" + }, + "outputs": [], + "source": [ + "import metrax\n", + "\n", + "# Define the metrics we want to calculate.\n", + "metrics_to_compute = {\n", + " 'Precision': metrax.Precision,\n", + " 'Recall': metrax.Recall,\n", + " 'AUCPR': metrax.AUCPR,\n", + " 'AUCROC': metrax.AUCROC,\n", + "}\n", + "\n", + "# Define which of these metrics should receive sample weights.\n", + "metrics_with_weights = {'AUCPR', 'AUCROC'}\n", + "\n", + "\n", + "# --- Method 1: Full-Batch Calculation ---\n", + "print(\"--- Method 1: Full-Batch Calculation (on all 32 samples) ---\")\n", + "full_batch_results = {}\n", + "for name, MetricClass in metrics_to_compute.items():\n", + " # Conditionally add sample_weights for supported metrics.\n", + " if name in metrics_with_weights:\n", + " metric_state = MetricClass.from_model_output(\n", + " predictions=predictions,\n", + " labels=labels,\n", + " sample_weights=sample_weights\n", + " )\n", + " else:\n", + " metric_state = MetricClass.from_model_output(\n", + " predictions=predictions,\n", + " labels=labels\n", + " )\n", + " full_batch_results[name] = metric_state.compute()\n", + " print(f\"{name}: {full_batch_results[name]}\")\n", + "\n", + "\n", + "# --- Method 2: Iterative Merging by Batch ---\n", + "print(\"\\n--- Method 2: Iterative Merging (4 batches of 8 samples) ---\")\n", + "iterative_metrics = {\n", + " name: MetricClass.empty() for name, MetricClass in metrics_to_compute.items()\n", + "}\n", + "\n", + "for labels_b, predictions_b, weights_b in zip(labels_batched, predictions_batched, sample_weights_batched):\n", + " for name, MetricClass in metrics_to_compute.items():\n", + " if name in metrics_with_weights:\n", + " current_metric_state = MetricClass.from_model_output(\n", + " predictions=predictions_b,\n", + " labels=labels_b,\n", + " sample_weights=weights_b\n", + " )\n", + " else:\n", + " current_metric_state = MetricClass.from_model_output(\n", + " predictions=predictions_b,\n", + " labels=labels_b\n", + " )\n", + " iterative_metrics[name] = iterative_metrics[name].merge(current_metric_state)\n", + "\n", + "iterative_results = {}\n", + "for name, metric_state in iterative_metrics.items():\n", + " iterative_results[name] = metric_state.compute()\n", + " print(f\"{name}: {iterative_results[name]}\")\n", + "\n", + "\n", + "# --- Verification ---\n", + "print(\"\\n--- Verification ---\")\n", + "for name in metrics_to_compute.keys():\n", + " assert np.allclose(full_batch_results[name], iterative_results[name])\n", + "\n", + "print(\"āœ… Success! Both methods produce identical results.\")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "uCez70DmcIEr" + }, + "source": [ + "## Metrax NNX API (Object-Oriented)\n", + "\n", + "For users who prefer an object-oriented style, `metrax.nnx` provides stateful metric objects. This can simplify the code for iterative updates, as you update a single object in place.\n", + "\n", + "Each `metrax.nnx` metric inherits the NNX [`metric`](http://shortn/_VyVVvvsQ00) class and provides the following APIs:\n", + "\n", + "* `metric = Metric()`: Creates a stateful metric object.\n", + "* `metric.update()`: Updates the metric's internal state with new data.\n", + "* `metric.compute()`: Computes the final value from the accumulated state.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "s4WucRk0O2rg" + }, + "outputs": [], + "source": [ + "import metrax.nnx\n", + "\n", + "# Define the nnx metrics we want to calculate.\n", + "metrics_to_compute_nnx = {\n", + " 'Precision': metrax.nnx.Precision,\n", + " 'Recall': metrax.nnx.Recall,\n", + " 'AUCPR': metrax.nnx.AUCPR,\n", + " 'AUCROC': metrax.nnx.AUCROC,\n", + "}\n", + "\n", + "# Define which metrics require which specific arguments.\n", + "metrics_with_threshold = {'Precision', 'Recall'}\n", + "\n", + "\n", + "# --- Method 1: Full-Batch Calculation (nnx) ---\n", + "print(\"--- Method 1: Full-Batch Calculation with nnx ---\")\n", + "full_batch_metrics_nnx = {\n", + " name: MetricClass() for name, MetricClass in metrics_to_compute_nnx.items()\n", + "}\n", + "for name, metric_obj in full_batch_metrics_nnx.items():\n", + " update_kwargs = {'predictions': predictions, 'labels': labels}\n", + " if name in metrics_with_weights:\n", + " update_kwargs['sample_weights'] = sample_weights\n", + " if name in metrics_with_threshold:\n", + " update_kwargs['threshold'] = 0.5\n", + " metric_obj.update(**update_kwargs)\n", + "\n", + "full_batch_results_nnx = {}\n", + "for name, metric_obj in full_batch_metrics_nnx.items():\n", + " full_batch_results_nnx[name] = metric_obj.compute()\n", + " print(f\"{name}: {full_batch_results_nnx[name]}\")\n", + "\n", + "\n", + "# --- Method 2: Iterative Updating by Batch (nnx) ---\n", + "print(\"\\n--- Method 2: Iterative Updating with nnx ---\")\n", + "iterative_metrics_nnx = {\n", + " name: MetricClass() for name, MetricClass in metrics_to_compute_nnx.items()\n", + "}\n", + "for labels_b, predictions_b, weights_b in zip(labels_batched, predictions_batched, sample_weights_batched):\n", + " for name, metric_obj in iterative_metrics_nnx.items():\n", + " update_kwargs = {'predictions': predictions_b, 'labels': labels_b}\n", + " if name in metrics_with_weights:\n", + " update_kwargs['sample_weights'] = weights_b\n", + " if name in metrics_with_threshold:\n", + " update_kwargs['threshold'] = 0.5\n", + " metric_obj.update(**update_kwargs)\n", + "\n", + "iterative_results_nnx = {}\n", + "for name, metric_obj in iterative_metrics_nnx.items():\n", + " iterative_results_nnx[name] = metric_obj.compute()\n", + " print(f\"{name}: {iterative_results_nnx[name]}\")\n", + "\n", + "\n", + "# --- Verification ---\n", + "print(\"\\n--- Verification ---\")\n", + "for name in metrics_to_compute_nnx.keys():\n", + " assert np.allclose(full_batch_results_nnx[name], iterative_results_nnx[name])\n", + "\n", + "print(\"āœ… Success! Both methods produce identical results using the nnx API.\")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "9EqpSpHp16SD" + }, + "source": [ + "## Scaling to Multiple Devices\n", + "\n", + "`metrax` is designed from the ground up to work seamlessly in distributed environments. Let's explore the two primary ways to scale metric computations in JAX.\n", + "\n", + "### Method 1: The `pmap` Approach (Simple Data Parallelism)\n", + "\n", + "The legacy way to scale computations across multiple devices in JAX is with `jax.pmap`. It's a powerful transformation designed for **data parallelism**, where you want to run the exact same program on different slices of data. `metrax`'s composable metrics are perfectly suited for this.\n", + "\n", + "`pmap` is more than just a parallel loop; it's a transformation that handles several complex steps for you automatically. When you apply `pmap` to a function like `metrax.AUCPR.from_model_output`, here's what happens:\n", + "\n", + "1. **Compile (`jit`):** First, `pmap` **JIT-compiles** your function into highly efficient machine code that is optimized for your hardware (like GPUs or TPUs). You don't need to add `@jax.jit` yourself; it's an inherent part of the `pmap` process.\n", + "\n", + "2. **Distribute (Shard):** It then automatically splits your input arrays along their first (batch) axis and sends each data \"shard\" to a different device.\n", + "\n", + "3. **Execute in Parallel:** The compiled function runs simultaneously on all devices, with each device processing its own shard of the data. This step produces a distributed `metrax` object, where each device holds the intermediate metric state calculated from its local data.\n", + "\n", + "4. **Aggregate (Reduce):** To get a final, globally correct result, you call the `.reduce()` method on the output from `pmap`. This performs the necessary cross-device communication (an `all_reduce` sum) to combine the intermediate states from all devices into one.\n", + "\n", + "---\n", + "\n", + "### Method 2: The `jit` and `Mesh` Approach (Advanced Parallelism)\n", + "\n", + "For more advanced control over distributed computation, JAX provides an explicit sharding mechanism using the `jax.sharding` API. This **SPMD (Single Program, Multiple Data)** approach is more powerful and flexible than `pmap` and is the standard for large-scale models.\n", + "\n", + "Instead of `pmap`'s automatic behavior, you take explicit control over each step of the process:\n", + "\n", + "1. **Define a `Mesh`**: You first create a logical grid of your physical devices and give names to the axes (e.g., `Mesh(jax.devices(), ('data',))`). This describes the topology you'll be working with.\n", + "\n", + "2. **Create a `Sharding` Rule**: You specify exactly how each dimension of your array should be mapped to the mesh's axes. This is done using `NamedSharding` and `PartitionSpec`. For data parallelism, you would shard the batch axis of your data across the `'data'` axis of your mesh.\n", + "\n", + "3. **Explicitly Place Data**: You use `jax.device_put` to apply this sharding rule to your data arrays. At this point, your JAX arrays are \"aware\" of how they are distributed across the physical hardware.\n", + "\n", + "4. **`jit`-Compile the Function**: You write a function that looks like a normal, single-device calculation and decorate it with `@jax.jit`. When JAX's compiler sees that the inputs to this function are sharded arrays, it automatically generates a distributed version of the code. It implicitly handles all the necessary cross-device communication, so **no explicit `.reduce()` call is needed.**\n", + "\n", + "This method provides the fine-grained control that is essential for more complex scenarios like model parallelism (sharding a model's weights) in addition to data parallelism." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "n3r37x3sMQbx" + }, + "outputs": [], + "source": [ + "import jax\n", + "import numpy as np\n", + "import metrax\n", + "from jax.sharding import Mesh, NamedSharding, PartitionSpec\n", + "\n", + "# This script assumes that the JAX environment is configured for 4 devices\n", + "# and that the data arrays `predictions`, `labels`, and `sample_weights`\n", + "# have been created in a previous cell.\n", + "\n", + "# --- 1. Metric Calculation Functions ---\n", + "\n", + "# Method 1: pmap (Simple Data Parallelism)\n", + "def calculate_aucpr_pmap(predictions, labels, sample_weights):\n", + " \"\"\"\n", + " Distributes data across devices using pmap and aggregates with .reduce().\n", + " \"\"\"\n", + " n_devices = jax.device_count()\n", + " # Reshape data to have a leading device dimension for pmap.\n", + " sharded_preds = predictions.reshape((n_devices, -1))\n", + " sharded_labels = labels.reshape((n_devices, -1))\n", + " sharded_weights = sample_weights.reshape((n_devices, -1))\n", + "\n", + " # pmap the metric calculation function across devices.\n", + " per_device_metric = jax.pmap(metrax.AUCPR.from_model_output)(\n", + " predictions=sharded_preds,\n", + " labels=sharded_labels,\n", + " sample_weights=sharded_weights\n", + " )\n", + " # reduce() combines the states from all devices into one.\n", + " return per_device_metric.reduce()\n", + "\n", + "\n", + "# Method 2: jit + Mesh (Advanced SPMD Parallelism)\n", + "def calculate_aucpr_mesh(predictions, labels, sample_weights):\n", + " \"\"\"\n", + " Explicitly shards data across a device Mesh and calculates with jit.\n", + " \"\"\"\n", + " # 1. Define the device mesh and sharding rule.\n", + " mesh = Mesh(jax.devices(), axis_names=('data',))\n", + " sharding_rule = NamedSharding(mesh, PartitionSpec('data'))\n", + "\n", + " # 2. Explicitly move and shard the data onto the mesh.\n", + " sharded_predictions = jax.device_put(predictions, sharding_rule)\n", + " sharded_labels = jax.device_put(labels, sharding_rule)\n", + " sharded_weights = jax.device_put(sample_weights, sharding_rule)\n", + "\n", + " # 3. Define the function to be JIT-compiled.\n", + " def _calculate(preds, labs, weights):\n", + " return metrax.AUCPR.from_model_output(\n", + " predictions=preds, labels=labs, sample_weights=weights)\n", + "\n", + " # 4. JIT-compile the function with explicit sharding annotations.\n", + " # - in_shardings: Specifies how each input array is expected to be sharded.\n", + " # - out_sharding: Specifies the desired sharding for the output.\n", + " # 'None' means the output should be replicated on all devices.\n", + " jitted_calculate = jax.jit(\n", + " _calculate,\n", + " in_shardings=(sharding_rule, sharding_rule, sharding_rule),\n", + " out_shardings=None\n", + " )\n", + "\n", + " # The result is already a globally correct metric state, replicated on all devices.\n", + " return jitted_calculate(sharded_predictions, sharded_labels, sharded_weights)\n", + "\n", + "\n", + "# Baseline: Single-Device (Direct)\n", + "@jax.jit\n", + "def calculate_aucpr_direct(predictions, labels, sample_weights):\n", + " \"\"\"Computes AUCPR on the entire dataset on a single device.\"\"\"\n", + " return metrax.AUCPR.from_model_output(\n", + " predictions=predictions,\n", + " labels=labels,\n", + " sample_weights=sample_weights\n", + " )\n", + "\n", + "\n", + "# --- 2. Execution and Verification ---\n", + "print(\"\\nRunning all three AUCPR calculation methods...\")\n", + "\n", + "# Execute each of the three methods.\n", + "state_pmap = calculate_aucpr_pmap(predictions, labels, sample_weights)\n", + "state_mesh = calculate_aucpr_mesh(predictions, labels, sample_weights)\n", + "state_direct = calculate_aucpr_direct(predictions, labels, sample_weights)\n", + "\n", + "# Compute the final values from the metric states.\n", + "result_pmap = state_pmap.compute()\n", + "result_mesh = state_mesh.compute()\n", + "result_direct = state_direct.compute()\n", + "\n", + "# Ensure all computations are finished before verifying.\n", + "result_pmap.block_until_ready()\n", + "result_mesh.block_until_ready()\n", + "result_direct.block_until_ready()\n", + "\n", + "# Verify that all results are numerically identical.\n", + "assert np.allclose(result_pmap, result_direct, rtol=1e-6)\n", + "assert np.allclose(result_mesh, result_direct, rtol=1e-6)\n", + "\n", + "\n", + "# --- 3. Display Results ---\n", + "print(\"\\n\" + \"=\"*60)\n", + "print(\" Comparison of Multi-Device AUCPR Calculations\")\n", + "print(\"=\"*60)\n", + "print(f\"{'Method':<35} {'AUCPR Value'}\")\n", + "print(\"-\" * 60)\n", + "print(f\"{'Method 1: pmap':<35} {result_pmap}\")\n", + "print(f\"{'Method 2: jit + Mesh':<35} {result_mesh}\")\n", + "print(f\"{'Baseline: Direct Single-Device':<35} {result_direct}\")\n", + "print(\"=\"*60)\n", + "print(\"\\nāœ… Verification successful: All three methods yield identical results.\")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "C3YWS1_x19DJ" + }, + "source": [ + "## šŸ§  Advanced Use: Multi-Host Environments\n", + "\n", + "For large-scale training (e.g., on TPU Pods), JAX uses multiple hosts (controllers), each managing multiple devices. In these scenarios, `metrax` integrates seamlessly with JAX's sharding capabilities provided by `jax.sharding` and `Mesh`.\n", + "\n", + "By `jit`-compiling your training step with the appropriate device mesh, you can calculate metrics across hundreds or thousands of devices without changing the core metric logic.\n", + "\n", + "For a detailed example of multi-controller and Pathways training, please see our [**example training script**](https://source.corp.google.com/piper///depot/google3/third_party/py/metrax/examples/xm_launch.py). The example demonstrates a multi-host setup with two hosts, each controlling four local devices for a total of eight(`viperlite=4x2`)." + ] + } + ], + "metadata": { + "colab": { + "last_runtime": { + "build_target": "//third_party/py/metrax/examples:metrax_colab", + "kind": "private" + }, + "private_outputs": true, + "provenance": [ + { + "file_id": "/piper/depot/google3/third_party/py/metrax/examples/metrax_example.ipynb", + "timestamp": 1749579890641 + }, + { + "file_id": "1DB9uscin86F2qtFOkm_kONZLTvYLYupj", + "timestamp": 1749511551887 + } + ] + }, + "kernelspec": { + "display_name": "Python 3", + "name": "python3" + }, + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 0 +} \ No newline at end of file