def __call__(
      self,
      tokens: pz.nx.NamedArray,
      *,
      token_positions: pz.nx.NamedArray | None = None,
      **side_inputs,
  ) -> pz.nx.NamedArray:
    """Scores log-probabilities for the given inputs.

    Args:
      tokens: Array of token IDs, as an integer named array with a "seq" axis
        and possibly batch axes. Usually starts with the beginning-of-sequence
        token.
      token_positions: Array of token positions, as an integer named array with
        a "seq" axis and possibly batch axes. Usually starts with 0. Inferred to
        start from 0 and increment along the "seq" axis if not provided.
      **side_inputs: Side inputs, which will be forwarded to the body.

    Returns:
      The final matrix of logits from the embedding decoding layer, which
      (in the normal configuration) will have axes "seq" and "vocabulary".
    """
    if token_positions is None:
      token_positions = pz.nx.arange("seq", tokens.named_shape["seq"])
    return self.body(tokens, token_positions=token_positions, **side_inputs)
Scores log-probabilities for the given inputs.

    Args:
      tokens: Array of token IDs, as an integer named array with a "seq" axis
        and possibly batch axes. Usually starts with the beginning-of-sequence
        token.
      token_positions: Array of token positions, as an integer named array with
        a "seq" axis and possibly batch axes. Usually starts with 0. Inferred to
        start from 0 and increment along the "seq" axis if not provided.
      **side_inputs: Side inputs, which will be forwarded to the body.

    Returns:
      The final matrix of logits from the embedding decoding layer, which
      (in the normal configuration) will have axes "seq" and "vocabulary".
    
penzai/models/transformer/model_parts.py