def METHOD_NAME(self):
    # Note we avoid validating because self.total_count can be zero.
    return torch.distributions.Gamma(
        concentration=self.total_count,
        rate=torch.exp(-self.logits),
        validate_args=False,
    )