KL Divergence of Two Bernoullis

Problem

Compute the KL divergence in bits between two Bernoulli distributions with success probabilities p and q. Terms with zero probability follow the convention 0 log 0 = 0, and the function assumes 0 < q < 1.

Hint

Sum the two contributions p*log2(p/q) and (1-p)*log2((1-p)/(1-q)), skipping zero-probability terms.

Test cases

case 1

in: [0.5,0.5]

expected: 0

case 2

in: [0.75,0.5]

expected: 0.18872187554086717

case 3

in: [0,0.5]

expected: 1

case 4

in: [1,0.25]

expected: 2

case 5

in: [0.25,0.75]

expected: 0.792481250360578

Starter codePython

import math
def bernoulli_kl(p, q):
    # Your code here
    pass

5 test cases · run your code against them in the editor