src/transformers/time_series_utils.py/NegativeBinomialOutput/domain_map
class NegativeBinomialOutput:    def domain_map(cls, total_count: torch.Tensor, logits: torch.Tensor):
        total_count = cls.squareplus(total_count)
        return total_count.squeeze(-1), logits.squeeze(-1)