def _head_info(config: LlamalikeTransformerConfig):
  """Computes query, key, and value head axes and einsum names."""
  if config.query_head_multiplier == 1:
    common_head_axes = {"heads": config.num_kv_heads}
    qkv_einsum = {"heads": "h"}
    query_only_head_axes = {}
    q_einsum = {}
  elif config.num_kv_heads == 1:
    common_head_axes = {}
    qkv_einsum = {}
    query_only_head_axes = {"query_heads": config.query_head_multiplier}
    q_einsum = {"query_heads": "h"}
  else:
    common_head_axes = {"head_groups": config.num_kv_heads}
    qkv_einsum = {"head_groups": "hg"}
    query_only_head_axes = {"query_heads": config.query_head_multiplier}
    q_einsum = {"query_heads": "hq"}
  return (common_head_axes, qkv_einsum, query_only_head_axes, q_einsum)
Computes query, key, and value head axes and einsum names.
penzai/models/transformer/variants/llamalike_common.py