def METHOD_NAME(self, sample_shape=torch.Size()):
    samples = self.sample(sample_shape)
    probs = self._categorical.probs  # cached via @lazy_property
    return samples + (probs - probs.detach())