Why is a distribution object made not subscriptable?

1

I need to take a subset of a distribution object, then
compare the subset with another subset of another distribution object.
Because the distribution object is not subcriptable (so we can not index directly the distribution object),
my approach so far is to index the member variable, eg probs,
then reconstruct the distribution object later.

Is there any better way for such use-case?
And why is the distribution object made not subscriptable?

m = torch.distributions.Categorical(probs=torch.rand(5,3))
subset = m.probs[0:2] # Okay
subset = m[0:2] # what I mean with directly indexing the object, but: TypeError: 'Categorical' object is not subscriptable
2

I want to ping this thread, because I think distributions module gains more popularity in time-series forecasting right now, and I also have faced similar issue recently. I would much rather subscript the distribution before performing .log_prob for loss, instead of plaing around with mocking and then masking the returned values risking to run into domain constraints and so on.
One use case is missing or “ignore” target value that needs to be handled before computing probs and then masked back, which is not nice.

I don’t see any reason why tensor-style indexing can’t be implemented in those objects? I will be willing to implement it if there is an interest to merge this functionality!