KL Chain Rule Gap

Problem

Verify the KL chain rule numerically. Given joints pxy and qxy, return the absolute residual between KL(p(x, y) || q(x, y)) and the chain decomposition KL(p(x) || q(x)) + sum_x p(x) * KL(p(y | x) || q(y | x)). All divergences are in bits and the residual should be within floating-point error of 0.

Hint

Build the conditionals by dividing each row by its marginal.

Test cases

case 1

in: [[[0.25,0.25],[0.25,0.25]],[[0.1,0.4],[0.4,0.1]]]

expected: 0

case 2

in: [[[0.5,0],[0,0.5]],[[0.25,0.25],[0.25,0.25]]]

expected: 0

case 3

in: [[[0.4,0.1],[0.2,0.3]],[[0.25,0.25],[0.25,0.25]]]

expected: 0

case 4

in: [[[0.3,0.2],[0.1,0.4]],[[0.2,0.3],[0.3,0.2]]]

expected: 0

Starter codePython

import math
def kl_chain_rule_gap(pxy, qxy):
    # Your code here
    pass

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