def get_loss_per_group(loss_per_label: Array, weight_per_group: Array,
                       group_start: int, group_end: int) -> Array:
  """Computes loss per sample within a group of labels."""
  loss_per_group = loss_per_label[:, group_start:group_end].sum(1)
  loss_per_group *= weight_per_group
  return loss_per_group
Computes loss per sample within a group of labels.
language/mentionmemory/tasks/ultra_fine_entity_typing_task.py