def show_token_array(
    tokens: pz.nx.NamedArrayBase | jax.Array | np.ndarray,
    vocab: SentencePieceProcessorLike,
    axis_name_order: Sequence[str] | None = None,
):
  """Renders an array of tokens.

  Args:
    tokens: An array of token IDs.
    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.

  Returns:
    A visualization of the token array.
  """
  [tokens], axis_name_order = _fixup_tokens_and_name_order(
      [tokens], axis_name_order
  )

  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"))
    subparts = []
    for i in range(tokens.named_shape[axis_name_order[-1]]):
      tok = int(tokens[{**ixs, axis_name_order[-1]: i}])
      if vocab.IsControl(tok):
        piece = pz.ts.with_color(vocab.IdToPiece(tok), "gray")
      else:
        piece = repr(vocab.IdToPiece(tok))
        if (piece.startswith("'") or piece.startswith('"')) and (
            piece.endswith("'") or piece.endswith('"')
        ):
          piece = piece[1:-1]
        else:
          piece = pz.ts.with_color(piece, "orange")
      subparts.append(
          pz.ts.inline(
              pz.ts.integer_digitbox(tok, label=f"{tok}"),
              " ",
              piece,
              wrap=False,
          )
      )
      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 tokens.

  Args:
    tokens: An array of token IDs.
    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.

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