def _fixup_tokens_and_name_order(
    token_arrays: list[pz.nx.NamedArrayBase | jax.Array | np.ndarray],
    axis_name_order: Sequence[str] | None = None,
) -> tuple[list[pz.nx.NamedArrayView], Sequence[str]]:
  """Helper function to preprocess token arrays and axis names."""
  namified_arrays = []
  common_shape = {}
  for token_array in token_arrays:
    if not isinstance(token_array, pz.nx.NamedArrayBase):
      token_array = pz.nx.wrap(
          jax.device_put(
              token_array, device=jax.local_devices(backend="cpu")[0]
          )
      )
    if token_array.positional_shape:
      token_array = token_array.tag(
          *(f"axis{i}" for i in range(len(token_array.positional_shape)))
      )
    token_array = jax.device_get(token_array)
    common_shape.update(token_array.named_shape)
    namified_arrays.append(token_array.as_namedarrayview())

  out_arrays = []
  shaped_dummy = pz.nx.zeros(common_shape)
  for token_array in namified_arrays:
    # Use nmap as an easy way to broadcast.
    out_arrays.append(pz.nx.nmap(lambda x, y: x)(token_array, shaped_dummy))

  if axis_name_order is None:
    base_order = list(out_arrays[0].named_shape.keys())
    if "batch" in base_order:
      base_order.remove("batch")
      base_order.insert(0, "batch")
    if "seq" in base_order:
      base_order.remove("seq")
      base_order.append("seq")
    axis_name_order = base_order

  axis_name_order = list(axis_name_order)

  if not out_arrays[0].named_shape:
    out_arrays = [token_array[{"seq": None}] for token_array in out_arrays]
    axis_name_order.append("seq")

  return out_arrays, axis_name_order
Helper function to preprocess token arrays and axis names.
penzai/toolshed/token_visualization.py