def histogram(self, tag, values, bins, step=None):
    """Saves histogram of values.

    Args:
      tag: str: label for this data
      values: ndarray: will be flattened by this routine
      bins: number of bins in histogram, or array of bins for np.histogram
      step: int: training step
    """
    if step is None:
      step = self._step
    else:
      self._step = step
    values = np.array(values)
    bins = np.array(bins)
    values = np.reshape(values, -1)
    counts, limits = np.histogram(values, bins=bins)
    # boundary logic
    cum_counts = np.cumsum(np.greater(counts, 0, dtype=np.int32))
    start, end = np.searchsorted(
        cum_counts, [0, cum_counts[-1] - 1], side='right')
    start, end = int(start), int(end) + 1
    counts = (
        counts[start -
               1:end] if start > 0 else np.concatenate([[0], counts[:end]]))
    limits = limits[start:end + 1]
    sum_sq = values.dot(values)
    histo = tf.compat.v1.HistogramProto(
        min=values.min(),
        max=values.max(),
        num=len(values),
        sum=values.sum(),
        sum_squares=sum_sq,
        bucket_limit=limits.tolist(),
        bucket=counts.tolist())
    summary = tf.compat.v1.Summary(
        value=[tf.compat.v1.Summary.Value(tag=tag, histo=histo)])
    self.add_summary(summary, step)
Saves histogram of values.

    Args:
      tag: str: label for this data
      values: ndarray: will be flattened by this routine
      bins: number of bins in histogram, or array of bins for np.histogram
      step: int: training step
    
trax/jaxboard.py