sample_entropy.py

import torch
import torch.nn
from torch.autograd import Variable

def pairwise_euclidean(samples):

    B = samples.size(0)

    samples_norm = samples.mul(samples).sum(1)
    samples_norm = samples_norm.expand(B, B)

    dist_mat = samples.mm(samples.t()).mul(-2) + \
        samples_norm.add(samples_norm.t())
    return dist_mat


def sample_entropy(samples):
    # Assume B x C input

    dist_mat = pairwise_euclidean(samples)

    # Get max and add it to diag
    m = dist_mat.max().detach()
    
    dist_mat_d = dist_mat + \
        Variable(torch.eye(dist_mat.size(0)).type_as(samples.data) * (m.data[0] + 1))

    entropy = (dist_mat_d.min(1)[0] + 1e-4).log().sum()

    entropy *= (samples.size(1) + 0.) / samples.size(0)

    return entropy