def plot(self, tag, mpl_plt, step=None, close_plot=True):
    """Saves matplotlib plot output to summary image.

    Args:
      tag: str: label for this data
      mpl_plt: matplotlib stateful pyplot object with prepared plotting state
      step: int: training step
      close_plot: bool: automatically closes plot
    """
    if step is None:
      step = self._step
    else:
      self._step = step
    fig = mpl_plt.get_current_fig_manager()
    img_w, img_h = fig.canvas.get_width_height()
    image_buf = io.BytesIO()
    mpl_plt.savefig(image_buf, format='png')
    image_summary = tf.compat.v1.Summary.Image(
        encoded_image_string=image_buf.getvalue(),
        colorspace=4,  # RGBA
        height=img_h,
        width=img_w)
    summary = tf.compat.v1.Summary(
        value=[tf.compat.v1.Summary.Value(tag=tag, image=image_summary)])
    self.add_summary(summary, step)
    if close_plot:
      mpl_plt.close()
Saves matplotlib plot output to summary image.

    Args:
      tag: str: label for this data
      mpl_plt: matplotlib stateful pyplot object with prepared plotting state
      step: int: training step
      close_plot: bool: automatically closes plot
    
trax/jaxboard.py