def get_weight_per_group(labels: Array, group_start: int,
                         group_end: int) -> Array:
  """Computes which samples have at least one labels within a group."""
  label_per_group_exists = labels[:, group_start:group_end].sum(1) > 0
  label_per_group_exists = label_per_group_exists.astype(jnp.float32)
  return label_per_group_exists
Computes which samples have at least one labels within a group.
language/mentionmemory/tasks/ultra_fine_entity_typing_task.py