import torch
from torch.special import digamma

def stable_beta_entropy(alpha, beta):
    ab = alpha + beta
    entropy = (
        - torch.lgamma(alpha) - torch.lgamma(beta)
        + torch.lgamma(ab)
        + (alpha - 1) * digamma(alpha)
        + (beta - 1) * digamma(beta)
        - (ab - 2) * digamma(ab)
    )
    return entropy.clamp(min=0.0)

if __name__ == '__main__':

    alpha = torch.tensor([1.1, 1.2, 1.3])
    beta = torch.tensor([1.3, 1.2, 1.1])

    # Pytorch --> *NEGATIVE* Entropy, does mathemaically not make sense
    dist = torch.distributions.Beta(alpha, beta)
    print(f"Pytorch Beta Entropy: {dist.entropy()}")  # should all be > 0

    # Custom Implementation --> *POSITIVE* Entropy
    print(f"Custom Beta Entropy: {stable_beta_entropy(alpha,beta)}")``

Pytorch Beta Entropy: tensor([-0.0206, -0.0108, -0.0206])
Custom Beta Entropy: tensor([0.0206, 0.0108, 0.0206])