Computing KL divergence for many distributions

Question

I have a matrix of test probability distributions:

qs = np.array([[0.1, 0.6], 
               [0.9, 0.4] ])

(sums up to 1 in each column) and "true" distribution:

p = np.array([0.5, 0.5])

I would like to calculate the KL divergence from p to every column of qs in TensorFlow . I know that there is a function tf.distributions.kl_divergence , but it takes just two distributions...

Answer 1

You could loop through the columns:

sess = tf.InteractiveSession()
A = tf.distributions.Categorical(probs=[0.5, 0.5])

for i in range(2):
    B = tf.distributions.Categorical(probs=list(qs[:,i]))
    print(tf.distributions.kl_divergence(A, B).eval())

Answer 2

This solution works for arbitrarily many columns and uses TensorFlow only:

qs = np.array([[0.9, 0.6, 0.4, 0.5],
               [0.1, 0.4, 0.6, 0.5]])
p = np.array([0.5, 0.5])
# Format `p` to a column vector, for constistency reason
p_t = p.reshape([-1, 1])

# Convert to categorical distributions. Transpose part is important
dist_qs = tf.distributions.Categorical(probs=tf.transpose(qs))
dist_p = tf.distributions.Categorical(probs=tf.transpose(p))

# Calculate KL divergences for qs and broadcasted p 
tf.distributions.kl_divergence(dist_p, dist_qs).eval()

This gives:

array([0.51082562, 0.020411  , 0.020411  , 0.])

what looks as a promising result.

Answer 3

Here's how to do it:

sess = tf.InteractiveSession()
A = tf.distributions.Categorical(probs=[0.1, 0.6])
B = tf.distributions.Categorical(probs=[0.9, 0.4])
print(tf.distributions.kl_divergence(A, B).eval())

Output: 0.7773577