Export Aggregated Metrics report to tensorboard.

PiperOrigin-RevId: 834280995

Author: mridul-sahu

checkpoint/orbax/checkpoint/_src/multihost/multislice.py

@@ -156,7 +156,7 @@ def get_device_memory() -> int:
       'NVIDIA H100 80GB HBM3': int(80e9),
       'NVIDIA H200': int(144e9),
       'NVIDIA B200': int(183e9),
-      'NVIDIA B300 SXM6 AC': int(275e9)
+      'NVIDIA B300 SXM6 AC': int(275e9),
   }
   memory = hbm_memory.get(device.device_kind, None)
   if memory is None:

checkpoint/orbax/checkpoint/_src/testing/benchmarks/core/metric.py

@@ -670,6 +670,28 @@ def export_to_tensorboard(self, tensorboard_dir: epath.Path):
                 )
             },
         )
+      # Write Aggreagated metrics as text
+      aggregated_metrics = []
+      aggregated_stats_dict, metric_units = self._aggregate_metrics(results)
+
+      if not aggregated_stats_dict:
+        aggregated_metrics.append("No successful runs to aggregate.")
+        continue
+
+      for key, stats in aggregated_stats_dict.items():
+        unit = metric_units[key]
+        aggregated_metrics.append(
+            f"{key}: {stats.mean:.4f} +/- {stats.std:.4f} {unit} (min:"
+            f" {stats.min:.4f}, max: {stats.max:.4f}, n={stats.count})"
+        )
+        aggregated_metrics_str = "\n".join(aggregated_metrics)
+        writer.write_texts(
+            step=0,
+            texts={
+                "aggregated_metrics": f"<pre>{aggregated_metrics_str}</pre>"
+            },
+        )
+
       writer.flush()
       writer.close()
     logging.info("Finished writing metrics to TensorBoard.")