Add histogram support with wandb-compatible API (#309)

Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: gradio-pr-bot <gradio-pr-bot@users.noreply.github.com>

Author: 3 people

.changeset/brave-suns-think.md

@@ -0,0 +1,5 @@
+---
+"trackio": minor
+---
+
+feat:Add histogram support with wandb-compatible API

examples/training-with-histograms.py

@@ -0,0 +1,37 @@
+import numpy as np
+
+import trackio
+
+run = trackio.init(project="histogram-training")
+
+num_epochs = 10
+batch_size = 32
+learning_rate = 0.001
+
+for epoch in range(num_epochs):
+    epoch_losses = []
+    epoch_weights = []
+
+    for batch in range(20):
+        loss = np.random.exponential(scale=2.0) * (1 - epoch * 0.05)
+        epoch_losses.append(loss)
+
+        weights = np.random.normal(0, 1 - epoch * 0.08, 1000)
+        epoch_weights.extend(weights)
+
+    avg_loss = np.mean(epoch_losses)
+
+    trackio.log(
+        {
+            "loss": avg_loss,
+            "learning_rate": learning_rate * (0.95**epoch),
+            "loss_distribution": trackio.Histogram(epoch_losses, num_bins=20),
+            "weight_distribution": trackio.Histogram(epoch_weights, num_bins=50),
+        },
+        step=epoch,
+    )
+
+gradients = np.random.laplace(0, 0.1, 5000)
+trackio.log({"final_gradients": trackio.Histogram(gradients, num_bins=30)})
+
+trackio.finish()

tests/e2e/test_histogram_e2e.py

@@ -0,0 +1,43 @@
+import numpy as np
+
+import trackio
+from trackio.sqlite_storage import SQLiteStorage
+
+
+def test_histogram_with_trackio_log(temp_dir):
+    """Test logging histograms with trackio."""
+    run = trackio.init(project="test_histogram")
+
+    data1 = np.random.randn(1000)
+    data2 = np.random.exponential(2, 500)
+
+    trackio.log(
+        {
+            "normal_dist": trackio.Histogram(data1),
+            "exp_dist": trackio.Histogram(data2, num_bins=30),
+        }
+    )
+
+    hist, bins = np.histogram(data1, bins=25)
+    trackio.log({"precomputed": trackio.Histogram(np_histogram=(hist, bins))})
+
+    trackio.finish()
+
+    logs = SQLiteStorage.get_logs("test_histogram", run.name)
+
+    assert len(logs) == 2
+
+    assert "normal_dist" in logs[0]
+    assert logs[0]["normal_dist"]["_type"] == "trackio.histogram"
+    assert len(logs[0]["normal_dist"]["bins"]) == 65
+    assert len(logs[0]["normal_dist"]["values"]) == 64
+
+    assert "exp_dist" in logs[0]
+    assert logs[0]["exp_dist"]["_type"] == "trackio.histogram"
+    assert len(logs[0]["exp_dist"]["bins"]) == 31
+    assert len(logs[0]["exp_dist"]["values"]) == 30
+
+    assert "precomputed" in logs[1]
+    assert logs[1]["precomputed"]["_type"] == "trackio.histogram"
+    assert len(logs[1]["precomputed"]["bins"]) == 26
+    assert len(logs[1]["precomputed"]["values"]) == 25

tests/test_histogram.py

@@ -0,0 +1,58 @@
+import numpy as np
+import pytest
+
+from trackio.histogram import Histogram
+
+
+def test_histogram_from_sequence():
+    """Test creating histogram from a sequence of values."""
+    data = [1, 2, 3, 4, 5, 5, 5, 6, 7, 8, 8, 9]
+    hist = Histogram(data)
+
+    assert hist.bins is not None
+    assert hist.histogram is not None
+    assert len(hist.bins) == 65
+    assert len(hist.histogram) == 64
+    assert sum(hist.histogram) == len(data)
+
+
+def test_histogram_from_np_histogram():
+    """Test creating histogram from pre-computed numpy histogram."""
+    data = np.random.randn(500)
+    np_hist, np_bins = np.histogram(data, bins=30)
+
+    hist = Histogram(np_histogram=(np_hist, np_bins))
+
+    assert np.array_equal(hist.bins, np_bins)
+    assert np.array_equal(hist.histogram, np_hist)
+    assert len(hist.bins) == 31
+    assert len(hist.histogram) == 30
+
+
+def test_histogram_to_dict():
+    """Test histogram serialization to dictionary."""
+    data = np.random.randn(100)
+    hist = Histogram(data, num_bins=10)
+
+    hist_dict = hist._to_dict()
+
+    assert hist_dict["_type"] == "trackio.histogram"
+    assert "bins" in hist_dict
+    assert "values" in hist_dict
+    assert isinstance(hist_dict["bins"], list)
+    assert isinstance(hist_dict["values"], list)
+    assert len(hist_dict["bins"]) == 11
+    assert len(hist_dict["values"]) == 10
+
+
+def test_histogram_invalid_inputs():
+    """Test histogram with invalid inputs."""
+    with pytest.raises(
+        ValueError, match="Must provide either sequence or np_histogram"
+    ):
+        Histogram()
+
+    with pytest.raises(
+        ValueError, match="Cannot provide both sequence and np_histogram"
+    ):
+        Histogram([1, 2, 3], np_histogram=([1, 2], [0, 1, 2]))

trackio/__init__.py

@@ -1,4 +1,5 @@
 import hashlib
+import json
 import logging
 import os
 import warnings
@@ -13,6 +14,7 @@
 from huggingface_hub import SpaceStorage
 
 from trackio import context_vars, deploy, utils
+from trackio.histogram import Histogram
 from trackio.imports import import_csv, import_tf_events
 from trackio.media import TrackioImage, TrackioVideo
 from trackio.run import Run
@@ -30,7 +32,9 @@
     module="gradio.helpers",
 )
 
-__version__ = Path(__file__).parent.joinpath("version.txt").read_text().strip()
+__version__ = json.loads(Path(__file__).parent.joinpath("package.json").read_text())[
+    "version"
+]
 
 __all__ = [
     "init",
@@ -42,6 +46,7 @@
     "Image",
     "Video",
     "Table",
+    "Histogram",
 ]
 
 Image = TrackioImage

trackio/histogram.py

@@ -0,0 +1,68 @@
+from typing import Any
+
+import numpy as np
+
+
+class Histogram:
+    """
+    Histogram data type for Trackio, compatible with wandb.Histogram.
+
+    Example:
+        ```python
+        import trackio
+        import numpy as np
+
+        # Create histogram from sequence
+        data = np.random.randn(1000)
+        trackio.log({"distribution": trackio.Histogram(data)})
+
+        # Create histogram from numpy histogram
+        hist, bins = np.histogram(data, bins=30)
+        trackio.log({"distribution": trackio.Histogram(np_histogram=(hist, bins))})
+
+        # Specify custom number of bins
+        trackio.log({"distribution": trackio.Histogram(data, num_bins=50)})
+        ```
+
+    Args:
+        sequence: Optional sequence of values to create histogram from
+        np_histogram: Optional pre-computed numpy histogram (hist, bins) tuple
+        num_bins: Number of bins for the histogram (default 64, max 512)
+    """
+
+    TYPE = "trackio.histogram"
+
+    def __init__(
+        self,
+        sequence: Any = None,
+        np_histogram: tuple | None = None,
+        num_bins: int = 64,
+    ):
+        if sequence is None and np_histogram is None:
+            raise ValueError("Must provide either sequence or np_histogram")
+
+        if sequence is not None and np_histogram is not None:
+            raise ValueError("Cannot provide both sequence and np_histogram")
+
+        num_bins = min(num_bins, 512)
+
+        if np_histogram is not None:
+            self.histogram, self.bins = np_histogram
+            self.histogram = np.asarray(self.histogram)
+            self.bins = np.asarray(self.bins)
+        else:
+            data = np.asarray(sequence).flatten()
+            data = data[np.isfinite(data)]
+            if len(data) == 0:
+                self.histogram = np.array([])
+                self.bins = np.array([])
+            else:
+                self.histogram, self.bins = np.histogram(data, bins=num_bins)
+
+    def _to_dict(self) -> dict:
+        """Convert histogram to dictionary for storage."""
+        return {
+            "_type": self.TYPE,
+            "bins": self.bins.tolist(),
+            "values": self.histogram.tolist(),
+        }

trackio/run.py

@@ -7,6 +7,7 @@
 from gradio_client import Client, handle_file
 
 from trackio import utils
+from trackio.histogram import Histogram
 from trackio.media import TrackioMedia
 from trackio.sqlite_storage import SQLiteStorage
 from trackio.table import Table
@@ -139,7 +140,7 @@ def _process_media(self, metrics, step: int | None) -> dict:
     @staticmethod
     def _replace_tables(metrics):
         for k, v in metrics.items():
-            if isinstance(v, Table):
+            if isinstance(v, (Table, Histogram)):
                 metrics[k] = v._to_dict()
 
     def log(self, metrics: dict, step: int | None = None):

trackio/ui/main.py

@@ -14,6 +14,7 @@
 try:
     import trackio.utils as utils
     from trackio.file_storage import FileStorage
+    from trackio.histogram import Histogram
     from trackio.media import TrackioImage, TrackioVideo
     from trackio.sqlite_storage import SQLiteStorage
     from trackio.table import Table
@@ -25,6 +26,7 @@
 except ImportError:
     import utils
     from file_storage import FileStorage
+    from histogram import Histogram
     from media import TrackioImage, TrackioVideo
     from sqlite_storage import SQLiteStorage
     from table import Table
@@ -1121,6 +1123,71 @@ def update_dashboard(
                                         f"Column {metric_name} failed to render as a table: {e}"
                                     )
 
+        # Display histograms
+        histogram_cols = set(master_df.columns) - {
+            "run",
+            "step",
+            "timestamp",
+            "data_type",
+        }
+        if metrics_subset:
+            histogram_cols = [c for c in histogram_cols if c in metrics_subset]
+        if metric_filter and metric_filter.strip():
+            histogram_cols = filter_metrics_by_regex(
+                list(histogram_cols), metric_filter
+            )
+
+        actual_histogram_count = sum(
+            1
+            for metric_name in histogram_cols
+            if not (metric_df := master_df.dropna(subset=[metric_name])).empty
+            and isinstance(value := metric_df[metric_name].iloc[-1], dict)
+            and value.get("_type") == Histogram.TYPE
+        )
+
+        if actual_histogram_count > 0:
+            with gr.Accordion(f"histograms ({actual_histogram_count})", open=True):
+                with gr.Row(key="histogram-row"):
+                    for metric_idx, metric_name in enumerate(histogram_cols):
+                        metric_df = master_df.dropna(subset=[metric_name])
+                        if not metric_df.empty:
+                            value = metric_df[metric_name].iloc[-1]
+                            if (
+                                isinstance(value, dict)
+                                and "_type" in value
+                                and value["_type"] == Histogram.TYPE
+                            ):
+                                try:
+                                    bins = value.get("bins", [])
+                                    values = value.get("values", [])
+
+                                    if len(bins) > 0 and len(values) > 0:
+                                        bin_centers = [
+                                            (bins[i] + bins[i + 1]) / 2
+                                            for i in range(len(bins) - 1)
+                                        ]
+
+                                        df = pd.DataFrame(
+                                            {"bin_center": bin_centers, "count": values}
+                                        )
+
+                                        gr.BarPlot(
+                                            df,
+                                            x="bin_center",
+                                            y="count",
+                                            title=f"{metric_name} (latest)",
+                                            x_title="Value",
+                                            y_title="Count",
+                                            key=f"histogram-{metric_idx}",
+                                            show_fullscreen_button=True,
+                                            min_width=400,
+                                            show_export_button=True,
+                                        )
+                                except Exception as e:
+                                    gr.Warning(
+                                        f"Column {metric_name} failed to render as a histogram: {e}"
+                                    )
+
     with grouped_runs_panel:
 
         @gr.render(