Avoid stream sync in softmax

[spinjo]

• When ptr and index are both provided in torch_geometric.utils.softmax, use index to set the output_size in torch.repeat_interleave
• In test/utils/test_softmax.py, add a line that tests softmax(src, ptr=ptr, index=index)

If the output_size argument is not provided in torch.repeat_interleave, then the output size is calculated on the fly, at the cost of a stream synchronization, see https://docs.pytorch.org/docs/stable/generated/torch.repeat_interleave.html. But if the torch_geometric.utils.softmax function gets the index argument in addition to ptr, we get output_size for free. So we should use it.