WalisonCruz/function-gemma

1from typing import Sequence2 3import numpy as np4 5 6class Histogram:7    """8    Histogram data type for Trackio, compatible with wandb.Histogram.9 10    Args:11        sequence (`np.ndarray` or `Sequence[float]` or `Sequence[int]`, *optional*):12            Sequence of values to create the histogram from.13        np_histogram (`tuple`, *optional*):14            Pre-computed NumPy histogram as a `(hist, bins)` tuple.15        num_bins (`int`, *optional*, defaults to `64`):16            Number of bins for the histogram (maximum `512`).17 18    Example:19        ```python20        import trackio21        import numpy as np22 23        # Create histogram from sequence24        data = np.random.randn(1000)25        trackio.log({"distribution": trackio.Histogram(data)})26 27        # Create histogram from numpy histogram28        hist, bins = np.histogram(data, bins=30)29        trackio.log({"distribution": trackio.Histogram(np_histogram=(hist, bins))})30 31        # Specify custom number of bins32        trackio.log({"distribution": trackio.Histogram(data, num_bins=50)})33        ```34    """35 36    TYPE = "trackio.histogram"37 38    def __init__(39        self,40        sequence: np.ndarray | Sequence[float] | Sequence[int] | None = None,41        np_histogram: tuple | None = None,42        num_bins: int = 64,43    ):44        if sequence is None and np_histogram is None:45            raise ValueError("Must provide either sequence or np_histogram")46 47        if sequence is not None and np_histogram is not None:48            raise ValueError("Cannot provide both sequence and np_histogram")49 50        num_bins = min(num_bins, 512)51 52        if np_histogram is not None:53            self.histogram, self.bins = np_histogram54            self.histogram = np.asarray(self.histogram)55            self.bins = np.asarray(self.bins)56        else:57            data = np.asarray(sequence).flatten()58            data = data[np.isfinite(data)]59            if len(data) == 0:60                self.histogram = np.array([])61                self.bins = np.array([])62            else:63                self.histogram, self.bins = np.histogram(data, bins=num_bins)64 65    def _to_dict(self) -> dict:66        """Convert histogram to dictionary for storage."""67        return {68            "_type": self.TYPE,69            "bins": self.bins.tolist(),70            "values": self.histogram.tolist(),71        }72