diff --git a/metrax_example.ipynb b/examples/metrax_example.ipynb similarity index 87% rename from metrax_example.ipynb rename to examples/metrax_example.ipynb index 6a6b55e..ca59610 100644 --- a/metrax_example.ipynb +++ b/examples/metrax_example.ipynb @@ -60,11 +60,11 @@ "# 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", + "\n", "print(\"Configuring JAX to simulate 4 CPU devices...\")\n", - "os.environ['XLA_FLAGS'] = '--xla_force_host_platform_device_count=4'\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", @@ -122,8 +122,12 @@ "# --- 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", + "predictions_batched = predictions.reshape(N_BATCHES, BATCH_SIZE).astype(\n", + " np.float32\n", + ")\n", + "sample_weights_batched = sample_weights.reshape(N_BATCHES, BATCH_SIZE).astype(\n", + " np.float32\n", + ")\n", "\n", "# --- 4. Data Shape Verification ---\n", "print(\"āœ… Data generation complete. Verifying array shapes:\")\n", @@ -215,8 +219,7 @@ "full_batch_results = {}\n", "for name, MetricClass in metrics_to_compute.items():\n", " metric_state = MetricClass.from_model_output(\n", - " predictions=predictions,\n", - " labels=labels\n", + " predictions=predictions, labels=labels\n", " )\n", " full_batch_results[name] = metric_state.compute()\n", " print(f\"{name}: {full_batch_results[name]}\")" @@ -233,16 +236,18 @@ "# --- Method 2: Iterative Merging by Batch (Unweighted) ---\n", "print(\"\\n--- Method 2: Iterative Merging (Unweighted) ---\")\n", "iterative_metrics = {\n", - " name: MetricClass.empty() for name, MetricClass in metrics_to_compute.items()\n", + " name: MetricClass.empty()\n", + " for name, MetricClass in metrics_to_compute.items()\n", "}\n", "\n", "for labels_b, predictions_b in zip(labels_batched, predictions_batched):\n", " for name, MetricClass in metrics_to_compute.items():\n", " current_metric_state = MetricClass.from_model_output(\n", - " predictions=predictions_b,\n", - " labels=labels_b\n", + " predictions=predictions_b, labels=labels_b\n", + " )\n", + " iterative_metrics[name] = iterative_metrics[name].merge(\n", + " current_metric_state\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", @@ -290,9 +295,7 @@ "full_batch_results_weighted = {}\n", "for name, MetricClass in metrics_with_weights.items():\n", " metric_state = MetricClass.from_model_output(\n", - " predictions=predictions,\n", - " labels=labels,\n", - " sample_weights=sample_weights\n", + " predictions=predictions, labels=labels, sample_weights=sample_weights\n", " )\n", " full_batch_results_weighted[name] = metric_state.compute()\n", " print(f\"{name}: {full_batch_results_weighted[name]}\")" @@ -309,16 +312,19 @@ "# --- Method 2: Iterative Merging by Batch (Weighted) ---\n", "print(\"\\n--- Method 2: Iterative Merging (Weighted) ---\")\n", "iterative_metrics_weighted = {\n", - " name: MetricClass.empty() for name, MetricClass in metrics_with_weights.items()\n", + " name: MetricClass.empty()\n", + " for name, MetricClass in metrics_with_weights.items()\n", "}\n", - "for labels_b, predictions_b, weights_b in zip(labels_batched, predictions_batched, sample_weights_batched):\n", + "for labels_b, predictions_b, weights_b in zip(\n", + " labels_batched, predictions_batched, sample_weights_batched\n", + "):\n", " for name, MetricClass in metrics_with_weights.items():\n", " current_metric_state = MetricClass.from_model_output(\n", - " predictions=predictions_b,\n", - " labels=labels_b,\n", - " sample_weights=weights_b\n", + " predictions=predictions_b, labels=labels_b, sample_weights=weights_b\n", + " )\n", + " iterative_metrics_weighted[name] = iterative_metrics_weighted[name].merge(\n", + " current_metric_state\n", " )\n", - " iterative_metrics_weighted[name] = iterative_metrics_weighted[name].merge(current_metric_state)\n", "\n", "iterative_results_weighted = {}\n", "for name, metric_state in iterative_metrics_weighted.items():\n", @@ -411,7 +417,9 @@ " name: MetricClass() for name, MetricClass in metrics_to_compute_nnx.items()\n", "}\n", "\n", - "for labels_b, predictions_b, _ in zip(labels_batched, predictions_batched, sample_weights_batched):\n", + "for labels_b, predictions_b, _ in zip(\n", + " labels_batched, predictions_batched, sample_weights_batched\n", + "):\n", " for name, metric_obj in iterative_metrics_nnx.items():\n", " metric_obj.update(predictions=predictions_b, labels=labels_b)\n", "\n", @@ -461,15 +469,14 @@ "# --- Method 1: Full-Batch Calculation with Sample Weights ---\n", "print(\"--- Method 1: Full-Batch Calculation with nnx (Weighted) ---\")\n", "full_batch_metrics_weighted = {\n", - " name: MetricClass() for name, MetricClass in weighted_metrics_to_compute_nnx.items()\n", + " name: MetricClass()\n", + " for name, MetricClass in weighted_metrics_to_compute_nnx.items()\n", "}\n", "\n", "for name, metric_obj in full_batch_metrics_weighted.items():\n", " # Update with predictions, labels, AND sample_weights\n", " metric_obj.update(\n", - " predictions=predictions,\n", - " labels=labels,\n", - " sample_weights=sample_weights\n", + " predictions=predictions, labels=labels, sample_weights=sample_weights\n", " )\n", "\n", "full_batch_results_weighted = {}\n", @@ -489,15 +496,16 @@ "# --- Method 2: Iterative Updating with Sample Weights ---\n", "print(\"\\n--- Method 2: Iterative Updating with nnx (Weighted) ---\")\n", "iterative_metrics_weighted = {\n", - " name: MetricClass() for name, MetricClass in weighted_metrics_to_compute_nnx.items()\n", + " name: MetricClass()\n", + " for name, MetricClass in weighted_metrics_to_compute_nnx.items()\n", "}\n", "\n", - "for labels_b, predictions_b, weights_b in zip(labels_batched, predictions_batched, sample_weights_batched):\n", + "for labels_b, predictions_b, weights_b in zip(\n", + " labels_batched, predictions_batched, sample_weights_batched\n", + "):\n", " for name, metric_obj in iterative_metrics_weighted.items():\n", " metric_obj.update(\n", - " predictions=predictions_b,\n", - " labels=labels_b,\n", - " sample_weights=weights_b\n", + " predictions=predictions_b, labels=labels_b, sample_weights=weights_b\n", " )\n", "\n", "iterative_results_weighted = {}\n", @@ -539,22 +547,21 @@ "outputs": [], "source": [ "import jax\n", - "import numpy as np\n", - "import metrax\n", "from jax.sharding import Mesh, NamedSharding, PartitionSpec\n", + "import metrax\n", + "import numpy as np\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", + "\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", + " predictions=predictions, labels=labels, sample_weights=sample_weights\n", " )" ] }, @@ -568,35 +575,34 @@ "source": [ "# Advanced SPMD Parallelism: jit + Mesh\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", + " \"\"\"Explicitly shards data across a device Mesh and calculates with jit.\"\"\"\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", "\n", - " # The result is already a globally correct metric state, replicated on all devices.\n", - " return jitted_calculate(sharded_predictions, sharded_labels, sharded_weights)" + " # 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)" ] }, { @@ -628,16 +634,18 @@ "assert np.allclose(result_pmap, result_direct, rtol=1e-6)\n", "assert np.allclose(result_mesh, result_direct, rtol=1e-6)\n", "\n", - "print(\"\\n\" + \"=\"*60)\n", + "print(\"\\n\" + \"=\" * 60)\n", "print(\" Comparison of Multi-Device AUCPR Calculations\")\n", - "print(\"=\"*60)\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.\")" + "print(\"=\" * 60)\n", + "print(\n", + " \"\\nāœ… Verification successful: All three methods yield identical results.\"\n", + ")" ] }, {