If you have only moderate tensor sizes, broadcasting probably is an easy way to achieve this:
import math
logits = torch.tensor([1.,2,1,3,1])
segment_id = torch.tensor([0,0,0,1,1])
segment_all_ids = torch.arange(2)
neginf = torch.full((), -math.inf, device=logits.device, dtype=logits.dtype)
logits_per_seg = torch.where(segment_id==segment_all_ids[:, None], logits, neginf)
p_per_seg = logits_per_seg.softmax(1)
p = p_per_seg.sum(0)

Best regards
Thomas