import tensorflow as tf

            sqrt_diag[:, n][:, None] if diag else sqrt[n, :, :][None, :, :], # 1 x M x M or M x 1
            K if shared_k else K_batch[n, :, :][None,:,:])) # 1 x M x M or M x M
    kl_sum =tf.reduce_sum(kl_sum)
    assert_almost_equal(kl_sum.eval(), kl_batch.eval())
