def show_token_scores(
    tokens: pz.nx.NamedArrayBase | jax.Array | np.ndarray,
    scores: pz.nx.NamedArrayBase | jax.Array | np.ndarray,
    vocab: SentencePieceProcessorLike,
    axis_name_order: Sequence[str] | None = None,
    vmax: float | None = None,
):
  """Renders an array of token scores, but with the scored tokens as words.

  Args:
    tokens: An array of token IDs.
    scores: An array of token scores of the same shape as `tokens`.
    vocab: Vocabulary, usually a SentencePiece tokenizer.
    axis_name_order: Optional ordering of the axis names. The last axis name
      will be laid out horizontally. Inferred if not given.
    vmax: Value at which to truncate the colormap.

  Returns:
    A visualization of the token array.
  """
  [tokens, scores], axis_name_order = _fixup_tokens_and_name_order(
      [tokens, scores], axis_name_order
  )
  if vmax is None:
    root_mean_square = jnp.sqrt(jnp.mean(jnp.square(scores.data_array)))
    vmax = 3 * root_mean_square
  parts = []
  index_iters = []
  for name in axis_name_order[:-1]:
    index_iters.append([(name, i) for i in range(tokens.named_shape[name])])
  for outer in itertools.product(*index_iters):
    ixs = dict(outer)
    if ixs:
      parts.append(pz.ts.styled(f"\n{repr(ixs)}:", "font-style: italic"))
    else:
      # Make space for hover tooltips.
      parts.append("\n")
    subparts = []
    for i in range(tokens.named_shape[axis_name_order[-1]]):
      tok = int(tokens[{**ixs, axis_name_order[-1]: i}])
      score = float(scores[{**ixs, axis_name_order[-1]: i}])
      subparts.append(
          pz.ts.text_on_color(text=vocab.IdToPiece(tok), value=score, vmax=vmax)
      )
      subparts.append("")

    # Add an indentation level, but allow sequences to wrap in the indented
    # block.
    parts.append(figures.indented(pz.ts.inline(*subparts, wrap=True)))

  return pz.ts.inline(*parts)
Renders an array of token scores, but with the scored tokens as words.

  Args:
    tokens: An array of token IDs.
    scores: An array of token scores of the same shape as `tokens`.
    vocab: Vocabulary, usually a SentencePiece tokenizer.
    axis_name_order: Optional ordering of the axis names. The last axis name
      will be laid out horizontally. Inferred if not given.
    vmax: Value at which to truncate the colormap.

  Returns:
    A visualization of the token array.
  
penzai/toolshed/token_visualization.py