Divergência KL no TensorFlow

Eu tenho dois tensores,prob_a eprob_b com forma[None, 1000], e eu quero calcular a divergência KL deprob_a paraprob_b. Existe uma função interna para isso no TensorFlow? Eu tentei usartf.contrib.distributions.kl(prob_a, prob_b) mas dá:

NotImplementedError: No KL(dist_a || dist_b) registered for dist_a type Tensor and dist_b type Tensor

Se não houver uma função interna, qual seria uma boa solução alternativa?