def text(self, tag, textdata, step=None):
    """Saves a text summary.

    Args:
      tag: str: label for this data
      textdata: string, or 1D/2D list/numpy array of strings
      step: int: training step
    Note: markdown formatting is rendered by tensorboard.
    """
    if step is None:
      step = self._step
    else:
      self._step = step
    smd = tf.compat.v1.SummaryMetadata(
        plugin_data=tf.compat.v1.SummaryMetadata.PluginData(plugin_name='text'))
    if isinstance(textdata, (str, bytes)):
      tensor = tf.make_tensor_proto(
          values=[textdata.encode(encoding='utf_8')], shape=(1,))
    else:
      textdata = np.array(textdata)  # convert lists, jax arrays, etc.
      datashape = np.shape(textdata)
      if len(datashape) == 1:
        tensor = tf.make_tensor_proto(
            values=[td.encode(encoding='utf_8') for td in textdata],
            shape=(datashape[0],))
      elif len(datashape) == 2:
        tensor = tf.make_tensor_proto(
            values=[
                td.encode(encoding='utf_8') for td in np.reshape(textdata, -1)
            ],
            shape=(datashape[0], datashape[1]))
    summary = tf.compat.v1.Summary(
        value=[tf.compat.v1.Summary.Value(
            tag=tag, metadata=smd, tensor=tensor)])
    self.add_summary(summary, step)
Saves a text summary.

    Args:
      tag: str: label for this data
      textdata: string, or 1D/2D list/numpy array of strings
      step: int: training step
    Note: markdown formatting is rendered by tensorboard.
    
trax/jaxboard.py