def METHOD_NAME(self, sample_shape=torch.Size()):
    if not isinstance(sample_shape, torch.Size):
        sample_shape = torch.Size(sample_shape)
    probs_2d = self.probs.reshape(-1, self._num_events)
    samples_2d = torch.multinomial(probs_2d, sample_shape.numel(), True).T
    return samples_2d.reshape(self._extended_shape(sample_shape))