Skip to content

Restructure examples directory for open-source sync. #158

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
copybara-service wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Restructure examples directory for open-source sync. #158

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
138 changes: 73 additions & 65 deletions metrax_example.ipynb → examples/metrax_example.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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]}\")"
Expand All @@ -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",
Expand Down Expand Up @@ -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]}\")"
Expand All @@ -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",
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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",
Expand All @@ -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",
Expand Down Expand Up @@ -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",
" )"
]
},
Expand All @@ -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)"
]
},
{
Expand Down Expand Up @@ -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",
")"
]
},
{
Expand Down
Loading