Is there any Softmax implementation with sections along the dim (blocky Softmax) in PyTorch?

Question

For example, given logits, dim, and boundary,

boundary = torch.tensor([[0, 3, 4, 8, 0]
                         [1, 3, 5, 7, 9]]

# representing sections look