path: root/tensorflow/contrib/bayesflow/python/kernel_tests/csiszar_divergence_test.py

# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for Csiszar Divergence Ops."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np

from tensorflow.contrib.bayesflow.python.ops import csiszar_divergence_impl
from tensorflow.contrib.distributions.python.ops import mvn_diag as mvn_diag_lib
from tensorflow.contrib.distributions.python.ops import mvn_full_covariance as mvn_full_lib
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import gradients_impl
from tensorflow.python.ops import linalg_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import nn_ops
from tensorflow.python.ops.distributions import kullback_leibler
from tensorflow.python.ops.distributions import normal as normal_lib
from tensorflow.python.platform import test


cd = csiszar_divergence_impl


def tridiag(d, diag_value, offdiag_value):
  """d x d matrix with given value on diag, and one super/sub diag."""
  diag_mat = linalg_ops.eye(d) * (diag_value - offdiag_value)
  three_bands = array_ops.matrix_band_part(
      array_ops.fill([d, d], offdiag_value), 1, 1)
  return diag_mat + three_bands


class AmariAlphaTest(test.TestCase):

  def setUp(self):
    self._logu = np.linspace(-10., 10, 100)
    self._u = np.exp(self._logu)

  def test_at_zero(self):
    for alpha in [-1., 0., 1., 2.]:
      for normalized in [True, False]:
        with self.test_session(graph=ops.Graph()):
          self.assertAllClose(
              cd.amari_alpha(0., alpha=alpha,
                             self_normalized=normalized).eval(),
              0.)

  def test_correct_when_alpha0(self):
    with self.test_session():
      self.assertAllClose(
          cd.amari_alpha(self._logu, alpha=0.).eval(),
          -self._logu)

      self.assertAllClose(
          cd.amari_alpha(self._logu, alpha=0., self_normalized=True).eval(),
          -self._logu + (self._u - 1.))

  def test_correct_when_alpha1(self):
    with self.test_session():
      self.assertAllClose(
          cd.amari_alpha(self._logu, alpha=1.).eval(),
          self._u * self._logu)

      self.assertAllClose(
          cd.amari_alpha(self._logu, alpha=1., self_normalized=True).eval(),
          self._u * self._logu - (self._u - 1.))

  def test_correct_when_alpha_not_01(self):
    for alpha in [-2, -1., -0.5, 0.5, 2.]:
      with self.test_session(graph=ops.Graph()):
        self.assertAllClose(
            cd.amari_alpha(self._logu,
                           alpha=alpha,
                           self_normalized=False).eval(),
            ((self._u**alpha - 1)) / (alpha * (alpha - 1.)))

        self.assertAllClose(
            cd.amari_alpha(self._logu,
                           alpha=alpha,
                           self_normalized=True).eval(),
            ((self._u**alpha - 1.)
             - alpha * (self._u - 1)) / (alpha * (alpha - 1.)))


class KLReverseTest(test.TestCase):

  def setUp(self):
    self._logu = np.linspace(-10., 10, 100)
    self._u = np.exp(self._logu)

  def test_at_zero(self):
    for normalized in [True, False]:
      with self.test_session(graph=ops.Graph()):
        self.assertAllClose(
            cd.kl_reverse(0., self_normalized=normalized).eval(),
            0.)

  def test_correct(self):
    with self.test_session():
      self.assertAllClose(
          cd.kl_reverse(self._logu).eval(),
          -self._logu)

      self.assertAllClose(
          cd.kl_reverse(self._logu, self_normalized=True).eval(),
          -self._logu + (self._u - 1.))


class KLForwardTest(test.TestCase):

  def setUp(self):
    self._logu = np.linspace(-10., 10, 100)
    self._u = np.exp(self._logu)

  def test_at_zero(self):
    for normalized in [True, False]:
      with self.test_session(graph=ops.Graph()):
        self.assertAllClose(
            cd.kl_forward(0., self_normalized=normalized).eval(),
            0.)

  def test_correct(self):
    with self.test_session():
      self.assertAllClose(
          cd.kl_forward(self._logu).eval(),
          self._u * self._logu)

      self.assertAllClose(
          cd.kl_forward(self._logu, self_normalized=True).eval(),
          self._u * self._logu - (self._u - 1.))


class JensenShannonTest(test.TestCase):

  def setUp(self):
    self._logu = np.linspace(-10., 10, 100)
    self._u = np.exp(self._logu)

  def test_at_zero(self):
    with self.test_session():
      self.assertAllClose(cd.jensen_shannon(0.).eval(), np.log(0.25))

  def test_symmetric(self):
    with self.test_session():
      self.assertAllClose(
          cd.jensen_shannon(self._logu).eval(),
          cd.symmetrized_csiszar_function(
              self._logu, cd.jensen_shannon).eval())

      self.assertAllClose(
          cd.jensen_shannon(self._logu, self_normalized=True).eval(),
          cd.symmetrized_csiszar_function(
              self._logu,
              lambda x: cd.jensen_shannon(x, self_normalized=True)).eval())

  def test_correct(self):
    with self.test_session():
      self.assertAllClose(
          cd.jensen_shannon(self._logu).eval(),
          (self._u * self._logu
           - (1 + self._u) * np.log1p(self._u)))

      self.assertAllClose(
          cd.jensen_shannon(self._logu, self_normalized=True).eval(),
          (self._u * self._logu
           - (1 + self._u) * np.log((1 + self._u) / 2)))


class ArithmeticGeometricMeanTest(test.TestCase):

  def setUp(self):
    self._logu = np.linspace(-10., 10, 100)
    self._u = np.exp(self._logu)

  def test_at_zero(self):
    with self.test_session():
      self.assertAllClose(cd.arithmetic_geometric(0.).eval(), np.log(4))
      self.assertAllClose(
          cd.arithmetic_geometric(0., self_normalized=True).eval(), 0.)

  def test_symmetric(self):
    with self.test_session():
      self.assertAllClose(
          cd.arithmetic_geometric(self._logu).eval(),
          cd.symmetrized_csiszar_function(
              self._logu, cd.arithmetic_geometric).eval())

  def test_correct(self):
    with self.test_session():
      self.assertAllClose(
          cd.arithmetic_geometric(self._logu).eval(),
          (1. + self._u) * np.log((1. + self._u) / np.sqrt(self._u)))

      self.assertAllClose(
          cd.arithmetic_geometric(self._logu, self_normalized=True).eval(),
          (1. + self._u) * np.log(0.5 * (1. + self._u) / np.sqrt(self._u)))


class TotalVariationTest(test.TestCase):

  def setUp(self):
    self._logu = np.linspace(-10., 10, 100)
    self._u = np.exp(self._logu)

  def test_at_zero(self):
    with self.test_session():
      self.assertAllClose(cd.total_variation(0.).eval(), 0.)

  def test_correct(self):
    with self.test_session():
      self.assertAllClose(
          cd.total_variation(self._logu).eval(),
          0.5 * np.abs(self._u - 1))


class PearsonTest(test.TestCase):

  def setUp(self):
    self._logu = np.linspace(-10., 10, 100)
    self._u = np.exp(self._logu)

  def test_at_zero(self):
    with self.test_session():
      self.assertAllClose(cd.pearson(0.).eval(), 0.)

  def test_correct(self):
    with self.test_session():
      self.assertAllClose(
          cd.pearson(self._logu).eval(),
          np.square(self._u - 1))


class SquaredHellingerTest(test.TestCase):

  def setUp(self):
    self._logu = np.linspace(-10., 10, 100)
    self._u = np.exp(self._logu)

  def test_at_zero(self):
    with self.test_session():
      self.assertAllClose(cd.squared_hellinger(0.).eval(), 0.)

  def test_symmetric(self):
    with self.test_session():
      self.assertAllClose(
          cd.squared_hellinger(self._logu).eval(),
          cd.symmetrized_csiszar_function(
              self._logu, cd.squared_hellinger).eval())

  def test_correct(self):
    with self.test_session():
      self.assertAllClose(
          cd.squared_hellinger(self._logu).eval(),
          np.square(np.sqrt(self._u) - 1))


class TriangularTest(test.TestCase):

  def setUp(self):
    self._logu = np.linspace(-10., 10, 100)
    self._u = np.exp(self._logu)

  def test_at_zero(self):
    with self.test_session():
      self.assertAllClose(cd.triangular(0.).eval(), 0.)

  def test_symmetric(self):
    with self.test_session():
      self.assertAllClose(
          cd.triangular(self._logu).eval(),
          cd.symmetrized_csiszar_function(
              self._logu, cd.triangular).eval())

  def test_correct(self):
    with self.test_session():
      self.assertAllClose(
          cd.triangular(self._logu).eval(),
          np.square(self._u - 1) / (1 + self._u))


class TPowerTest(test.TestCase):

  def setUp(self):
    self._logu = np.linspace(-10., 10, 100)
    self._u = np.exp(self._logu)

  def test_at_zero(self):
    with self.test_session():
      self.assertAllClose(cd.t_power(0., t=-0.1).eval(), 0.)
      self.assertAllClose(cd.t_power(0., t=0.5).eval(), 0.)
      self.assertAllClose(cd.t_power(0., t=1.1).eval(), 0.)
      self.assertAllClose(
          cd.t_power(0., t=-0.1, self_normalized=True).eval(), 0.)
      self.assertAllClose(
          cd.t_power(0., t=0.5, self_normalized=True).eval(), 0.)
      self.assertAllClose(
          cd.t_power(0., t=1.1, self_normalized=True).eval(), 0.)

  def test_correct(self):
    with self.test_session():
      self.assertAllClose(
          cd.t_power(self._logu, t=np.float64(-0.1)).eval(),
          self._u ** -0.1 - 1.)
      self.assertAllClose(
          cd.t_power(self._logu, t=np.float64(0.5)).eval(),
          -self._u ** 0.5 + 1.)
      self.assertAllClose(
          cd.t_power(self._logu, t=np.float64(1.1)).eval(),
          self._u ** 1.1 - 1.)

  def test_correct_self_normalized(self):
    with self.test_session():
      self.assertAllClose(
          cd.t_power(self._logu, t=np.float64(-0.1),
                     self_normalized=True).eval(),
          self._u ** -0.1 - 1. + 0.1 * (self._u - 1.))
      self.assertAllClose(
          cd.t_power(self._logu, t=np.float64(0.5),
                     self_normalized=True).eval(),
          -self._u ** 0.5 + 1. + 0.5 * (self._u - 1.))
      self.assertAllClose(
          cd.t_power(self._logu, t=np.float64(1.1),
                     self_normalized=True).eval(),
          self._u ** 1.1 - 1. - 1.1 * (self._u - 1.))


class Log1pAbsTest(test.TestCase):

  def setUp(self):
    self._logu = np.linspace(-10., 10, 100)
    self._u = np.exp(self._logu)

  def test_at_zero(self):
    with self.test_session():
      self.assertAllClose(cd.log1p_abs(0.).eval(), 0.)

  def test_correct(self):
    with self.test_session():
      self.assertAllClose(
          cd.log1p_abs(self._logu).eval(),
          self._u**(np.sign(self._u - 1)) - 1)


class JeffreysTest(test.TestCase):

  def setUp(self):
    self._logu = np.linspace(-10., 10, 100)
    self._u = np.exp(self._logu)

  def test_at_zero(self):
    with self.test_session():
      self.assertAllClose(cd.jeffreys(0.).eval(), 0.)

  def test_symmetric(self):
    with self.test_session():
      self.assertAllClose(
          cd.jeffreys(self._logu).eval(),
          cd.symmetrized_csiszar_function(
              self._logu, cd.jeffreys).eval())

  def test_correct(self):
    with self.test_session():
      self.assertAllClose(
          cd.jeffreys(self._logu).eval(),
          0.5 * (self._u * self._logu - self._logu))


class ChiSquareTest(test.TestCase):

  def setUp(self):
    self._logu = np.linspace(-10., 10, 100)
    self._u = np.exp(self._logu)

  def test_at_zero(self):
    with self.test_session():
      self.assertAllClose(cd.chi_square(0.).eval(), 0.)

  def test_correct(self):
    with self.test_session():
      self.assertAllClose(
          cd.chi_square(self._logu).eval(),
          self._u**2 - 1)


class ModifiedGanTest(test.TestCase):

  def setUp(self):
    self._logu = np.linspace(-10., 10, 100)
    self._u = np.exp(self._logu)

  def test_at_zero(self):
    with self.test_session():
      self.assertAllClose(
          cd.modified_gan(0.).eval(), np.log(2))
      self.assertAllClose(
          cd.modified_gan(0., self_normalized=True).eval(), np.log(2))

  def test_correct(self):
    with self.test_session():
      self.assertAllClose(
          cd.modified_gan(self._logu).eval(),
          np.log1p(self._u) - self._logu)

      self.assertAllClose(
          cd.modified_gan(self._logu, self_normalized=True).eval(),
          np.log1p(self._u) - self._logu + 0.5 * (self._u - 1))


class SymmetrizedCsiszarFunctionTest(test.TestCase):

  def setUp(self):
    self._logu = np.linspace(-10., 10., 100)
    self._u = np.exp(self._logu)

  def test_jensen_shannon(self):
    with self.test_session():

      # The following functions come from the claim made in the
      # symmetrized_csiszar_function docstring.
      def js1(logu):
        return (-logu
                - (1. + math_ops.exp(logu)) * (
                    nn_ops.softplus(logu)))

      def js2(logu):
        return 2. * (math_ops.exp(logu) * (
            logu - nn_ops.softplus(logu)))

      self.assertAllClose(
          cd.symmetrized_csiszar_function(self._logu, js1).eval(),
          cd.jensen_shannon(self._logu).eval())

      self.assertAllClose(
          cd.symmetrized_csiszar_function(self._logu, js2).eval(),
          cd.jensen_shannon(self._logu).eval())

  def test_jeffreys(self):
    with self.test_session():
      self.assertAllClose(
          cd.symmetrized_csiszar_function(self._logu, cd.kl_reverse).eval(),
          cd.jeffreys(self._logu).eval())

      self.assertAllClose(
          cd.symmetrized_csiszar_function(self._logu, cd.kl_forward).eval(),
          cd.jeffreys(self._logu).eval())


class DualCsiszarFunctionTest(test.TestCase):

  def setUp(self):
    self._logu = np.linspace(-10., 10., 100)
    self._u = np.exp(self._logu)

  def test_kl_forward(self):
    with self.test_session():
      self.assertAllClose(
          cd.dual_csiszar_function(self._logu, cd.kl_forward).eval(),
          cd.kl_reverse(self._logu).eval())

  def test_kl_reverse(self):
    with self.test_session():
      self.assertAllClose(
          cd.dual_csiszar_function(self._logu, cd.kl_reverse).eval(),
          cd.kl_forward(self._logu).eval())


class MonteCarloCsiszarFDivergenceTest(test.TestCase):

  def test_kl_forward(self):
    with self.test_session() as sess:
      q = normal_lib.Normal(
          loc=np.ones(6),
          scale=np.array([0.5, 1.0, 1.5, 2.0, 2.5, 3.0]))

      p = normal_lib.Normal(loc=q.loc + 0.1, scale=q.scale - 0.2)

      approx_kl = cd.monte_carlo_csiszar_f_divergence(
          f=cd.kl_forward,
          p_log_prob=p.log_prob,
          q=q,
          num_draws=int(1e5),
          seed=1)

      approx_kl_self_normalized = cd.monte_carlo_csiszar_f_divergence(
          f=lambda logu: cd.kl_forward(logu, self_normalized=True),
          p_log_prob=p.log_prob,
          q=q,
          num_draws=int(1e5),
          seed=1)

      exact_kl = kullback_leibler.kl_divergence(p, q)

      [approx_kl_, approx_kl_self_normalized_, exact_kl_] = sess.run([
          approx_kl, approx_kl_self_normalized, exact_kl])

      self.assertAllClose(approx_kl_, exact_kl_,
                          rtol=0.08, atol=0.)

      self.assertAllClose(approx_kl_self_normalized_, exact_kl_,
                          rtol=0.02, atol=0.)

  def test_kl_reverse(self):
    with self.test_session() as sess:

      q = normal_lib.Normal(
          loc=np.ones(6),
          scale=np.array([0.5, 1.0, 1.5, 2.0, 2.5, 3.0]))

      p = normal_lib.Normal(loc=q.loc + 0.1, scale=q.scale - 0.2)

      approx_kl = cd.monte_carlo_csiszar_f_divergence(
          f=cd.kl_reverse,
          p_log_prob=p.log_prob,
          q=q,
          num_draws=int(1e5),
          seed=1)

      approx_kl_self_normalized = cd.monte_carlo_csiszar_f_divergence(
          f=lambda logu: cd.kl_reverse(logu, self_normalized=True),
          p_log_prob=p.log_prob,
          q=q,
          num_draws=int(1e5),
          seed=1)

      exact_kl = kullback_leibler.kl_divergence(q, p)

      [approx_kl_, approx_kl_self_normalized_, exact_kl_] = sess.run([
          approx_kl, approx_kl_self_normalized, exact_kl])

      self.assertAllClose(approx_kl_, exact_kl_,
                          rtol=0.07, atol=0.)

      self.assertAllClose(approx_kl_self_normalized_, exact_kl_,
                          rtol=0.02, atol=0.)

  def test_kl_reverse_multidim(self):

    with self.test_session() as sess:
      d = 5  # Dimension

      p = mvn_full_lib.MultivariateNormalFullCovariance(
          covariance_matrix=tridiag(d, diag_value=1, offdiag_value=0.5))

      q = mvn_diag_lib.MultivariateNormalDiag(scale_diag=[0.5]*d)

      approx_kl = cd.monte_carlo_csiszar_f_divergence(
          f=cd.kl_reverse,
          p_log_prob=p.log_prob,
          q=q,
          num_draws=int(1e5),
          seed=1)

      approx_kl_self_normalized = cd.monte_carlo_csiszar_f_divergence(
          f=lambda logu: cd.kl_reverse(logu, self_normalized=True),
          p_log_prob=p.log_prob,
          q=q,
          num_draws=int(1e5),
          seed=1)

      exact_kl = kullback_leibler.kl_divergence(q, p)

      [approx_kl_, approx_kl_self_normalized_, exact_kl_] = sess.run([
          approx_kl, approx_kl_self_normalized, exact_kl])

      self.assertAllClose(approx_kl_, exact_kl_,
                          rtol=0.02, atol=0.)

      self.assertAllClose(approx_kl_self_normalized_, exact_kl_,
                          rtol=0.08, atol=0.)

  def test_kl_forward_multidim(self):

    with self.test_session() as sess:
      d = 5  # Dimension

      p = mvn_full_lib.MultivariateNormalFullCovariance(
          covariance_matrix=tridiag(d, diag_value=1, offdiag_value=0.5))

      # Variance is very high when approximating Forward KL, so we make
      # scale_diag larger than in test_kl_reverse_multidim. This ensures q
      # "covers" p and thus Var_q[p/q] is smaller.
      q = mvn_diag_lib.MultivariateNormalDiag(scale_diag=[1.]*d)

      approx_kl = cd.monte_carlo_csiszar_f_divergence(
          f=cd.kl_forward,
          p_log_prob=p.log_prob,
          q=q,
          num_draws=int(1e5),
          seed=1)

      approx_kl_self_normalized = cd.monte_carlo_csiszar_f_divergence(
          f=lambda logu: cd.kl_forward(logu, self_normalized=True),
          p_log_prob=p.log_prob,
          q=q,
          num_draws=int(1e5),
          seed=1)

      exact_kl = kullback_leibler.kl_divergence(p, q)

      [approx_kl_, approx_kl_self_normalized_, exact_kl_] = sess.run([
          approx_kl, approx_kl_self_normalized, exact_kl])

      self.assertAllClose(approx_kl_, exact_kl_,
                          rtol=0.06, atol=0.)

      self.assertAllClose(approx_kl_self_normalized_, exact_kl_,
                          rtol=0.05, atol=0.)

  def test_score_trick(self):

    with self.test_session() as sess:
      d = 5  # Dimension
      num_draws = int(1e5)
      seed = 1

      p = mvn_full_lib.MultivariateNormalFullCovariance(
          covariance_matrix=tridiag(d, diag_value=1, offdiag_value=0.5))

      # Variance is very high when approximating Forward KL, so we make
      # scale_diag larger than in test_kl_reverse_multidim. This ensures q
      # "covers" p and thus Var_q[p/q] is smaller.
      s = array_ops.constant(1.)
      q = mvn_diag_lib.MultivariateNormalDiag(
          scale_diag=array_ops.tile([s], [d]))

      approx_kl = cd.monte_carlo_csiszar_f_divergence(
          f=cd.kl_reverse,
          p_log_prob=p.log_prob,
          q=q,
          num_draws=num_draws,
          seed=seed)

      approx_kl_self_normalized = cd.monte_carlo_csiszar_f_divergence(
          f=lambda logu: cd.kl_reverse(logu, self_normalized=True),
          p_log_prob=p.log_prob,
          q=q,
          num_draws=num_draws,
          seed=seed)

      approx_kl_score_trick = cd.monte_carlo_csiszar_f_divergence(
          f=cd.kl_reverse,
          p_log_prob=p.log_prob,
          q=q,
          num_draws=num_draws,
          use_reparametrization=False,
          seed=seed)

      approx_kl_self_normalized_score_trick = (
          cd.monte_carlo_csiszar_f_divergence(
              f=lambda logu: cd.kl_reverse(logu, self_normalized=True),
              p_log_prob=p.log_prob,
              q=q,
              num_draws=num_draws,
              use_reparametrization=False,
              seed=seed))

      exact_kl = kullback_leibler.kl_divergence(q, p)

      grad_sum = lambda fs: gradients_impl.gradients(fs, s)[0]

      [
          approx_kl_grad_,
          approx_kl_self_normalized_grad_,
          approx_kl_score_trick_grad_,
          approx_kl_self_normalized_score_trick_grad_,
          exact_kl_grad_,
          approx_kl_,
          approx_kl_self_normalized_,
          approx_kl_score_trick_,
          approx_kl_self_normalized_score_trick_,
          exact_kl_,
      ] = sess.run([
          grad_sum(approx_kl),
          grad_sum(approx_kl_self_normalized),
          grad_sum(approx_kl_score_trick),
          grad_sum(approx_kl_self_normalized_score_trick),
          grad_sum(exact_kl),
          approx_kl,
          approx_kl_self_normalized,
          approx_kl_score_trick,
          approx_kl_self_normalized_score_trick,
          exact_kl,
      ])

      # Test average divergence.
      self.assertAllClose(approx_kl_, exact_kl_,
                          rtol=0.02, atol=0.)

      self.assertAllClose(approx_kl_self_normalized_, exact_kl_,
                          rtol=0.08, atol=0.)

      self.assertAllClose(approx_kl_score_trick_, exact_kl_,
                          rtol=0.02, atol=0.)

      self.assertAllClose(approx_kl_self_normalized_score_trick_, exact_kl_,
                          rtol=0.08, atol=0.)

      # Test average gradient-divergence.
      self.assertAllClose(approx_kl_grad_, exact_kl_grad_,
                          rtol=0.007, atol=0.)

      self.assertAllClose(approx_kl_self_normalized_grad_, exact_kl_grad_,
                          rtol=0.011, atol=0.)

      self.assertAllClose(approx_kl_score_trick_grad_, exact_kl_grad_,
                          rtol=0.018, atol=0.)

      self.assertAllClose(
          approx_kl_self_normalized_score_trick_grad_, exact_kl_grad_,
          rtol=0.017, atol=0.)


class CsiszarVIMCOTest(test.TestCase):

  def _csiszar_vimco_helper(self, logu):
    """Numpy implementation of `csiszar_vimco_helper`."""

    # Since this is a naive/intuitive implementation, we compensate by using the
    # highest precision we can.
    logu = np.float128(logu)
    n = logu.shape[0]
    u = np.exp(logu)
    loogeoavg_u = []  # Leave-one-out geometric-average of exp(logu).
    for j in range(n):
      loogeoavg_u.append(np.exp(np.mean(
          [logu[i, ...] for i in range(n) if i != j],
          axis=0)))
    loogeoavg_u = np.array(loogeoavg_u)

    loosum_u = []  # Leave-one-out sum of exp(logu).
    for j in range(n):
      loosum_u.append(np.sum(
          [u[i, ...] for i in range(n) if i != j],
          axis=0))
    loosum_u = np.array(loosum_u)

    # Natural log of the average u except each is swapped-out for its
    # leave-`i`-th-out Geometric average.
    log_sooavg_u = np.log(loosum_u + loogeoavg_u) - np.log(n)

    log_avg_u = np.log(np.mean(u, axis=0))
    return log_avg_u, log_sooavg_u

  def _csiszar_vimco_helper_grad(self, logu, delta):
    """Finite difference approximation of `grad(csiszar_vimco_helper, logu)`."""

    # This code actually estimates the sum of the Jacobiab because that's what
    # TF's `gradients` does.
    np_log_avg_u1, np_log_sooavg_u1 = self._csiszar_vimco_helper(
        logu[..., None] + np.diag([delta]*len(logu)))
    np_log_avg_u, np_log_sooavg_u = self._csiszar_vimco_helper(
        logu[..., None])
    return [
        (np_log_avg_u1 - np_log_avg_u) / delta,
        np.sum(np_log_sooavg_u1 - np_log_sooavg_u, axis=0) / delta,
    ]

  def test_vimco_helper_1(self):
    """Tests that function calculation correctly handles batches."""

    logu = np.linspace(-100., 100., 100).reshape([10, 2, 5])
    with self.test_session() as sess:
      np_log_avg_u, np_log_sooavg_u = self._csiszar_vimco_helper(logu)
      [log_avg_u, log_sooavg_u] = sess.run(cd.csiszar_vimco_helper(logu))
      self.assertAllClose(np_log_avg_u, log_avg_u,
                          rtol=1e-8, atol=0.)
      self.assertAllClose(np_log_sooavg_u, log_sooavg_u,
                          rtol=1e-8, atol=0.)

  def test_vimco_helper_2(self):
    """Tests that function calculation correctly handles overflow."""

    # Using 700 (rather than 1e3) since naive numpy version can't handle higher.
    logu = np.float32([0., 700, -1, 1])
    with self.test_session() as sess:
      np_log_avg_u, np_log_sooavg_u = self._csiszar_vimco_helper(logu)
      [log_avg_u, log_sooavg_u] = sess.run(cd.csiszar_vimco_helper(logu))
      self.assertAllClose(np_log_avg_u, log_avg_u,
                          rtol=1e-6, atol=0.)
      self.assertAllClose(np_log_sooavg_u, log_sooavg_u,
                          rtol=1e-5, atol=0.)

  def test_vimco_helper_3(self):
    """Tests that function calculation correctly handles underlow."""

    logu = np.float32([0., -1000, -1, 1])
    with self.test_session() as sess:
      np_log_avg_u, np_log_sooavg_u = self._csiszar_vimco_helper(logu)
      [log_avg_u, log_sooavg_u] = sess.run(cd.csiszar_vimco_helper(logu))
      self.assertAllClose(np_log_avg_u, log_avg_u,
                          rtol=1e-5, atol=0.)
      self.assertAllClose(np_log_sooavg_u, log_sooavg_u,
                          rtol=1e-4, atol=1e-15)

  def test_vimco_helper_gradient_using_finite_difference_1(self):
    """Tests that gradient calculation correctly handles batches."""

    logu_ = np.linspace(-100., 100., 100).reshape([10, 2, 5])
    with self.test_session() as sess:
      logu = array_ops.constant(logu_)

      grad = lambda flogu: gradients_impl.gradients(flogu, logu)[0]
      log_avg_u, log_sooavg_u = cd.csiszar_vimco_helper(logu)

      [
          grad_log_avg_u,
          grad_log_sooavg_u,
      ] = sess.run([grad(log_avg_u), grad(log_sooavg_u)])

      # We skip checking against finite-difference approximation since it
      # doesn't support batches.

      # Verify claim in docstring.
      self.assertAllClose(
          np.ones_like(grad_log_avg_u.sum(axis=0)),
          grad_log_avg_u.sum(axis=0))
      self.assertAllClose(
          np.ones_like(grad_log_sooavg_u.mean(axis=0)),
          grad_log_sooavg_u.mean(axis=0))

  def test_vimco_helper_gradient_using_finite_difference_2(self):
    """Tests that gradient calculation correctly handles overflow."""

    delta = 1e-3
    logu_ = np.float32([0., 1000, -1, 1])
    with self.test_session() as sess:
      logu = array_ops.constant(logu_)

      [
          np_grad_log_avg_u,
          np_grad_log_sooavg_u,
      ] = self._csiszar_vimco_helper_grad(logu_, delta)

      grad = lambda flogu: gradients_impl.gradients(flogu, logu)[0]
      log_avg_u, log_sooavg_u = cd.csiszar_vimco_helper(logu)

      [
          grad_log_avg_u,
          grad_log_sooavg_u,
      ] = sess.run([grad(log_avg_u), grad(log_sooavg_u)])

      self.assertAllClose(np_grad_log_avg_u, grad_log_avg_u,
                          rtol=delta, atol=0.)
      self.assertAllClose(np_grad_log_sooavg_u, grad_log_sooavg_u,
                          rtol=delta, atol=0.)
      # Verify claim in docstring.
      self.assertAllClose(
          np.ones_like(grad_log_avg_u.sum(axis=0)),
          grad_log_avg_u.sum(axis=0))
      self.assertAllClose(
          np.ones_like(grad_log_sooavg_u.mean(axis=0)),
          grad_log_sooavg_u.mean(axis=0))

  def test_vimco_helper_gradient_using_finite_difference_3(self):
    """Tests that gradient calculation correctly handles underlow."""

    delta = 1e-3
    logu_ = np.float32([0., -1000, -1, 1])
    with self.test_session() as sess:
      logu = array_ops.constant(logu_)

      [
          np_grad_log_avg_u,
          np_grad_log_sooavg_u,
      ] = self._csiszar_vimco_helper_grad(logu_, delta)

      grad = lambda flogu: gradients_impl.gradients(flogu, logu)[0]
      log_avg_u, log_sooavg_u = cd.csiszar_vimco_helper(logu)

      [
          grad_log_avg_u,
          grad_log_sooavg_u,
      ] = sess.run([grad(log_avg_u), grad(log_sooavg_u)])

      self.assertAllClose(np_grad_log_avg_u, grad_log_avg_u,
                          rtol=delta, atol=0.)
      self.assertAllClose(np_grad_log_sooavg_u, grad_log_sooavg_u,
                          rtol=delta, atol=0.)
      # Verify claim in docstring.
      self.assertAllClose(
          np.ones_like(grad_log_avg_u.sum(axis=0)),
          grad_log_avg_u.sum(axis=0))
      self.assertAllClose(
          np.ones_like(grad_log_sooavg_u.mean(axis=0)),
          grad_log_sooavg_u.mean(axis=0))

  def test_vimco_and_gradient(self):

    with self.test_session() as sess:
      dims = 5  # Dimension
      num_draws = int(20)
      num_batch_draws = int(3)
      seed = 1

      f = lambda logu: cd.kl_reverse(logu, self_normalized=False)
      np_f = lambda logu: -logu

      p = mvn_full_lib.MultivariateNormalFullCovariance(
          covariance_matrix=tridiag(dims, diag_value=1, offdiag_value=0.5))

      # Variance is very high when approximating Forward KL, so we make
      # scale_diag larger than in test_kl_reverse_multidim. This ensures q
      # "covers" p and thus Var_q[p/q] is smaller.
      s = array_ops.constant(1.)
      q = mvn_diag_lib.MultivariateNormalDiag(
          scale_diag=array_ops.tile([s], [dims]))

      vimco = cd.csiszar_vimco(
          f=f,
          p_log_prob=p.log_prob,
          q=q,
          num_draws=num_draws,
          num_batch_draws=num_batch_draws,
          seed=seed)

      x = q.sample(sample_shape=[num_draws, num_batch_draws],
                   seed=seed)
      x = array_ops.stop_gradient(x)
      logu = p.log_prob(x) - q.log_prob(x)
      f_log_sum_u = f(cd.csiszar_vimco_helper(logu)[0])

      grad_sum = lambda fs: gradients_impl.gradients(fs, s)[0]

      def jacobian(x):
        # Warning: this function is slow and may not even finish if prod(shape)
        # is larger than, say, 100.
        shape = x.shape.as_list()
        assert all(s is not None for s in shape)
        x = array_ops.reshape(x, shape=[-1])
        r = [grad_sum(x[i]) for i in range(np.prod(shape))]
        return array_ops.reshape(array_ops.stack(r), shape=shape)

      [
          logu_,
          jacobian_logqx_,
          vimco_,
          grad_vimco_,
          f_log_sum_u_,
          grad_mean_f_log_sum_u_,
      ] = sess.run([
          logu,
          jacobian(q.log_prob(x)),
          vimco,
          grad_sum(vimco),
          f_log_sum_u,
          grad_sum(f_log_sum_u) / num_batch_draws,
      ])

      np_log_avg_u, np_log_sooavg_u = self._csiszar_vimco_helper(logu_)

      # Test VIMCO loss is correct.
      self.assertAllClose(np_f(np_log_avg_u).mean(axis=0), vimco_,
                          rtol=1e-5, atol=0.)

      # Test gradient of VIMCO loss is correct.
      #
      # To make this computation we'll inject two gradients from TF:
      # - grad[mean(f(log(sum(p(x)/q(x)))))]
      # - jacobian[log(q(x))].
      #
      # We now justify why using these (and only these) TF values for
      # ground-truth does not undermine the completeness of this test.
      #
      # Regarding `grad_mean_f_log_sum_u_`, note that we validate the
      # correctness of the zero-th order derivative (for each batch member).
      # Since `cd.csiszar_vimco_helper` itself does not manipulate any gradient
      # information, we can safely rely on TF.
      self.assertAllClose(np_f(np_log_avg_u), f_log_sum_u_, rtol=1e-4, atol=0.)
      #
      # Regarding `jacobian_logqx_`, note that testing the gradient of
      # `q.log_prob` is outside the scope of this unit-test thus we may safely
      # use TF to find it.

      # The `mean` is across batches and the `sum` is across iid samples.
      np_grad_vimco = (
          grad_mean_f_log_sum_u_
          + np.mean(
              np.sum(
                  jacobian_logqx_ * (np_f(np_log_avg_u)
                                     - np_f(np_log_sooavg_u)),
                  axis=0),
              axis=0))

      self.assertAllClose(np_grad_vimco, grad_vimco_,
                          rtol=1e-5, atol=0.)


if __name__ == "__main__":
  test.main()