<filename>tensorflow_probability/python/bijectors/glow_test.py
# Copyright 2020 The TensorFlow Probability Authors.
#
# 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 tensorflow_probability.python.bijectors.glow."""

import functools
# Dependency imports
import numpy as np
import tensorflow.compat.v2 as tf
import tensorflow_probability as tfp
from tensorflow_probability.python.bijectors import bijector_test_util
from tensorflow_probability.python.bijectors import blockwise
from tensorflow_probability.python.bijectors import composition
from tensorflow_probability.python.internal import test_util
from tensorflow_probability.python.math.gradient import batch_jacobian

tfb = tfp.bijectors
tfd = tfp.distributions


@test_util.test_all_tf_execution_regimes
class GlowTest(test_util.TestCase):

  def setUp(self):
    super(GlowTest, self).setUp()

    self.batch_shape = 5
    self.output_shape = [16, 16, 3]
    self.minval = -10.
    self.maxval = 10.

  def _create_glow(self, actnorm=False):

    glow_net = functools.partial(tfb.GlowDefaultNetwork, num_hidden=32)
    glow_exit = tfb.GlowDefaultExitNetwork
    glow_bijector = tfb.Glow(
        output_shape=self.output_shape,
        num_glow_blocks=2,
        num_steps_per_block=1,
        coupling_bijector_fn=glow_net,
        exit_bijector_fn=glow_exit,
        grab_after_block=[0.5, 0.5],
        use_actnorm=actnorm,
        seed=tfp.util.SeedStream(seed=42, salt='glow'))
    sigm = tfb.sigmoid.Sigmoid(
        low=self.minval, high=self.maxval)
    output = tfb.chain.Chain(
        [sigm, glow_bijector])
    return output

  def _make_images(self):
    sc = np.float32(self.maxval-self.minval)*0.999
    sh = np.float32(self.minval)+0.0005
    these_images = np.random.random([self.batch_shape]+self.output_shape)*sc+sh
    return np.float32(these_images)

  def testBijector(self):
    """Verify that the bijector is properly invertible."""
    bijection = self._create_glow()
    self.assertStartsWith(bijection.bijectors[1].name, 'glow')
    self.evaluate([v.initializer for v in bijection.variables])
    x = self._make_images()

    x = tf.constant(x, tf.float32)
    z = bijection.inverse(x)
    xprime = bijection.forward(tf.identity(z))
    zprime = bijection.inverse(xprime)

    # Absolute errors are < 1e-4, and need to use absolute since dealing with
    # values surrounding 0.
    self.assertAllClose(x, xprime, rtol=0., atol=1.e-4)
    self.assertAllClose(z, zprime, rtol=0., atol=1.e-4)
    self.assertAllClose(
        self.evaluate(-bijection.inverse_log_det_jacobian(x, 3)),
        self.evaluate(bijection.forward_log_det_jacobian(z, 1)),
        rtol=1e-4,
        atol=0.)

  def testBijectiveAndFinite(self):
    """Test that the bijector is invertible and that outputs are finite."""
    bijection = self._create_glow()
    self.evaluate([v.initializer for v in bijection.variables])
    ims = np.zeros([10]+self.output_shape)
    ims += np.linspace(self.minval+0.0001,
                       self.maxval-0.0001,
                       10, dtype=np.float32).reshape([10, 1, 1, 1])

    zs = np.float32(np.zeros([10, np.prod(self.output_shape)]))
    zs += np.linspace(0., 10., 10,
                      dtype=np.float32).reshape(10, 1)  # 10-sigma is far away.

    bijector_test_util.assert_bijective_and_finite(
        bijection, np.float32(zs), np.float32(ims), event_ndims=1,
        eval_func=self.evaluate, inverse_event_ndims=3,
        rtol=1e-3)

  def testDataInit_inverse(self):
    """Test that actnorm data dependent initialization works on inverse pass."""
    bijection = self._create_glow(actnorm=True)
    self.evaluate([v.initializer for v in bijection.variables])
    x = self._make_images()
    nblocks = 0
    made_a_check = False
    splits = bijection.bijectors[1].blockwise_splits
    splits = [[bs[0]+bs[1], bs[2]] for bs in splits]

    x = bijection.bijectors[0].inverse(x)
    for b in bijection.bijectors[1].bijectors:
      if isinstance(b, blockwise._Blockwise):
        x1, x2 = tf.split(x, splits[-2-nblocks], axis=-1)

        for bb in b.bijectors[0].bijectors:
          if isinstance(bb, tfb.glow.ActivationNormalization):
            x1 = self.evaluate(bb.inverse(x1))
            mean = self.evaluate(tf.reduce_mean(x1, axis=(-4, -3, -2)))
            stddev = self.evaluate(tf.math.reduce_std(x1, axis=(-4, -3, -2)))
            self.assertAllClose(mean, tf.zeros_like(mean), atol=1e-5)
            self.assertAllClose(stddev, tf.ones_like(stddev), atol=0, rtol=5e-5)
            made_a_check = True
          else:
            x1 = self.evaluate(bb.inverse(x1))
        x = tf.concat([x1, x2], axis=-1)
        nblocks += 1

      elif isinstance(b, composition.Composition):
        for bb in b.bijectors:
          x = self.evaluate(bb.inverse(x))
          if isinstance(bb, tfb.glow.ActivationNormalization):
            mean = tf.reduce_mean(x, axis=(-4, -3, -2))
            stddev = tf.math.reduce_std(x, axis=(-4, -3, -2))
            self.assertAllClose(mean, tf.zeros_like(mean), atol=1e-5)
            self.assertAllClose(stddev, tf.ones_like(stddev), atol=0, rtol=5e-5)
            made_a_check = True
      else:
        x = self.evaluate(b.inverse(x))

    assert made_a_check

  def testDataInit_forward(self):
    """Test that actnorm data dependent initialization works on forward pass."""
    bijection = self._create_glow(actnorm=True)
    self.evaluate([v.initializer for v in bijection.variables])
    y = np.float32(np.random.normal(0., 1., [5, 768]))
    nblocks = 0
    made_a_check = False
    splits = bijection.bijectors[1].blockwise_splits
    splits = [[bs[0]+bs[1], bs[2]] for bs in splits]

    for b in reversed(bijection.bijectors[1].bijectors):
      if isinstance(b, blockwise._Blockwise):
        y1, y2 = tf.split(y, splits[nblocks], axis=-1)

        for bb in reversed(b.bijectors[0].bijectors):
          if isinstance(bb, tfb.glow.ActivationNormalization):
            y1 = self.evaluate(bb.forward(y1))
            mean = self.evaluate(tf.reduce_mean(y1, axis=(-4, -3, -2)))
            stddev = self.evaluate(tf.math.reduce_std(y1, axis=(-4, -3, -2)))
            self.assertAllClose(mean, tf.zeros_like(mean), atol=1e-5)
            self.assertAllClose(stddev, tf.ones_like(stddev), atol=0, rtol=5e-5)
            made_a_check = True
          else:
            y1 = self.evaluate(bb.forward(y1))
        y = tf.concat([y1, y2], axis=-1)
        nblocks += 1

      elif isinstance(b, composition.Composition):
        for bb in reversed(b.bijectors):
          y = self.evaluate(bb.forward(y))
          if isinstance(bb, tfb.glow.ActivationNormalization):
            mean = tf.reduce_mean(y, axis=(-4, -3, -2))
            stddev = tf.math.reduce_std(y, axis=(-4, -3, -2))
            self.assertAllClose(mean, tf.zeros_like(mean), atol=1e-5)
            self.assertAllClose(stddev, tf.ones_like(stddev), atol=0, rtol=5e-5)
            made_a_check = True
      else:
        y = self.evaluate(b.forward(y))
    assert made_a_check

  def testInverseLogDetJacobian(self):
    """Test if log-det-jacobian agrees with numerical computation."""
    if not tf.executing_eagerly():
      self.skipTest('Theres a problem with numerical computation of Jacobian.'
                    'the bijector jacobian implementation  still returns'
                    'roughly the same values as it does in eager mode, so I'
                    'think our computation works here.')

    x = self._make_images()
    x = tf.constant(x, tf.float32)

    bijection = self._create_glow(actnorm=True)

    self.evaluate([v.initializer for v in bijection.variables])
    jacob_manual = self.evaluate(batch_jacobian(bijection.inverse, x))
    _, ldj_manual = np.linalg.slogdet(jacob_manual.reshape([5, 768, 768]))

    jacob = self.evaluate(bijection.inverse_log_det_jacobian(x, 3))
    self.assertAllClose(ldj_manual, jacob, rtol=1.e-5)

  def testBadBijectorFn(self):
    """Test a bad bijector function.

    Ensures that an error happens if the bijector outputs a tensor whose shape
    does not give a straightforward way to produce a bijector.
    """
    def bad_glow_net(input_shape):
      del input_shape
      return tfb.Pad()

    x = self._make_images()
    x = tf.constant(x, tf.float32)

    bijection = tfb.Glow(
        output_shape=self.output_shape,
        num_glow_blocks=2,
        num_steps_per_block=1,
        coupling_bijector_fn=bad_glow_net,
        exit_bijector_fn=tfb.GlowDefaultExitNetwork,
        grab_after_block=[0.5, 0.5],
        use_actnorm=False)
    self.evaluate([v.initializer for v in bijection.variables])
    try:
      z = bijection.inverse(x)
      del z
      raise Exception('If youre reading this, then something is wrong.')
    except ValueError:
      pass

  def testOtherBadBijectorFn(self):
    """Test another bad bijector function.

    This tests that an error is thrown if the bijector function produces
    something other than a bijector or a tensor
    """
    def other_bad_glow_net(input_shape):
      del input_shape
      return lambda x: tfd.Normal(x, 0)

    x = self._make_images()
    x = tf.constant(x, tf.float32)

    bijection = tfb.Glow(
        output_shape=self.output_shape,
        num_glow_blocks=2,
        num_steps_per_block=1,
        coupling_bijector_fn=other_bad_glow_net,
        exit_bijector_fn=tfb.GlowDefaultExitNetwork,
        grab_after_block=[0.5, 0.5],
        use_actnorm=False)
    self.evaluate([v.initializer for v in bijection.variables])
    try:
      z = bijection.inverse(x)
      del z
      raise Exception('If youre reading this, then something is wrong.')

    except ValueError:
      pass

  def testMultiDimensionalSampling(self):
    """This tests that the model can take samples in variable dimensions."""
    def dummy_glow_net(input_shape):
      del input_shape
      return lambda x: tfb.Identity()

    def dummy_exit_net(input_shape, output_chan):
      del input_shape, output_chan
      return lambda x: tfb.Identity()

    bijection = tfb.Glow(
        output_shape=self.output_shape,
        num_glow_blocks=2,
        num_steps_per_block=1,
        coupling_bijector_fn=dummy_glow_net,
        exit_bijector_fn=dummy_exit_net,
        grab_after_block=[0.5, 0.5],
        use_actnorm=False)
    dist = bijection(tfd.Independent(
        tfd.Normal(tf.zeros([16*16*3]), tf.ones(16*16*3)),
        reinterpreted_batch_ndims=1))

    single_samp2 = dist.sample(1)
    single_samp3 = dist.sample([1, 1])
    single_samp = dist.sample()
    self.assertShapeEqual(np.zeros([16, 16, 3]), single_samp)
    self.assertShapeEqual(np.zeros([1, 16, 16, 3]), single_samp2)
    self.assertShapeEqual(np.zeros([1, 1, 16, 16, 3]), single_samp3)

  def testMultiDimensionalInput(self):
    """This tests if the model runs with different batch shapes."""
    def dummy_glow_net(input_shape):
      del input_shape
      return lambda x: tfb.Identity()

    def dummy_exit_net(input_shape, output_chan):
      del input_shape, output_chan
      return lambda x: tfb.Identity()

    bijection = tfb.Glow(
        output_shape=self.output_shape,
        num_glow_blocks=2,
        num_steps_per_block=1,
        coupling_bijector_fn=dummy_glow_net,
        exit_bijector_fn=dummy_exit_net,
        grab_after_block=[0.5, 0.5],
        use_actnorm=False)

    single_z = bijection.inverse(tf.zeros([16, 16, 3]))
    single_z2 = bijection.inverse(tf.zeros([1, 16, 16, 3]))
    single_z3 = bijection.inverse(tf.zeros([1, 1, 16, 16, 3]))
    self.assertShapeEqual(np.zeros([16 * 16 * 3]), single_z)
    self.assertShapeEqual(np.zeros([1, 16 * 16 * 3]), single_z2)
    self.assertShapeEqual(np.zeros([1, 1, 16 * 16 * 3]), single_z3)

  def testDtypes(self):
    """Test if the bijector identifies dtype correctly."""
    ims = self._make_images()
    bijection = self._create_glow(False)
    z = bijection.inverse(ims)

    def float64_net(input_shape):
      input_nchan = input_shape[-1]
      return tf.keras.Sequential([
          tf.keras.layers.Input(input_shape, dtype=tf.float64),
          tf.keras.layers.Conv2D(
              2 * input_nchan, 3, padding='same', dtype=tf.float64)])
    def float64_exit(input_shape, output_chan):
      return tf.keras.Sequential([
          tf.keras.layers.Input(input_shape, dtype=tf.float64),
          tf.keras.layers.Conv2D(
              2*output_chan, 3, padding='same', dtype=tf.float64)])

    float64_bijection = tfb.Glow(
        output_shape=self.output_shape,
        num_glow_blocks=2,
        num_steps_per_block=1,
        coupling_bijector_fn=float64_net,
        exit_bijector_fn=float64_exit,
        grab_after_block=[0.5, 0.5]
        )
    zf64 = float64_bijection.inverse(tf.cast(ims, tf.float64))
    self.evaluate([v.initializer for v in bijection.variables])
    self.evaluate([v.initializer for v in float64_bijection.variables])
    self.assertAllFinite(self.evaluate(z))
    self.assertAllFinite(self.evaluate(zf64))

  def testBijectorFn(self):
    """Test if the bijector function works for additive coupling."""
    ims = self._make_images()
    def shiftfn(input_shape):
      input_nchan = input_shape[-1]
      return tf.keras.Sequential([
          tf.keras.layers.Input(input_shape),
          tf.keras.layers.Conv2D(
              input_nchan, 3, padding='same')])

    def shiftexitfn(input_shape, output_chan):
      return tf.keras.Sequential([
          tf.keras.layers.Input(input_shape),
          tf.keras.layers.Conv2D(
              output_chan, 3, padding='same')])

    shiftonlyglow = tfb.Glow(
        output_shape=self.output_shape,
        num_glow_blocks=2,
        num_steps_per_block=1,
        coupling_bijector_fn=shiftfn,
        exit_bijector_fn=shiftexitfn,
        grab_after_block=[0.5, 0.5]
        )
    z = shiftonlyglow.inverse(ims)
    self.evaluate([v.initializer for v in shiftonlyglow.variables])
    self.assertAllFinite(self.evaluate(z))


if __name__ == '__main__':
  test_util.main()

<filename>tensorflow_probability/python/bijectors/glow_test.py
# Copyright 2020 The TensorFlow Probability Authors.
#
# 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 tensorflow_probability.python.bijectors.glow."""

import functools
# Dependency imports
import numpy as np
import tensorflow.compat.v2 as tf
import tensorflow_probability as tfp
from tensorflow_probability.python.bijectors import bijector_test_util
from tensorflow_probability.python.bijectors import blockwise
from tensorflow_probability.python.bijectors import composition
from tensorflow_probability.python.internal import test_util
from tensorflow_probability.python.math.gradient import batch_jacobian

tfb = tfp.bijectors
tfd = tfp.distributions


@test_util.test_all_tf_execution_regimes
class GlowTest(test_util.TestCase):

  def setUp(self):
    super(GlowTest, self).setUp()

    self.batch_shape = 5
    self.output_shape = [16, 16, 3]
    self.minval = -10.
    self.maxval = 10.

  def _create_glow(self, actnorm=False):

    glow_net = functools.partial(tfb.GlowDefaultNetwork, num_hidden=32)
    glow_exit = tfb.GlowDefaultExitNetwork
    glow_bijector = tfb.Glow(
        output_shape=self.output_shape,
        num_glow_blocks=2,
        num_steps_per_block=1,
        coupling_bijector_fn=glow_net,
        exit_bijector_fn=glow_exit,
        grab_after_block=[0.5, 0.5],
        use_actnorm=actnorm,
        seed=tfp.util.SeedStream(seed=42, salt='glow'))
    sigm = tfb.sigmoid.Sigmoid(
        low=self.minval, high=self.maxval)
    output = tfb.chain.Chain(
        [sigm, glow_bijector])
    return output

  def _make_images(self):
    sc = np.float32(self.maxval-self.minval)*0.999
    sh = np.float32(self.minval)+0.0005
    these_images = np.random.random([self.batch_shape]+self.output_shape)*sc+sh
    return np.float32(these_images)

  def testBijector(self):
    """Verify that the bijector is properly invertible."""
    bijection = self._create_glow()
    self.assertStartsWith(bijection.bijectors[1].name, 'glow')
    self.evaluate([v.initializer for v in bijection.variables])
    x = self._make_images()

    x = tf.constant(x, tf.float32)
    z = bijection.inverse(x)
    xprime = bijection.forward(tf.identity(z))
    zprime = bijection.inverse(xprime)

    # Absolute errors are < 1e-4, and need to use absolute since dealing with
    # values surrounding 0.
    self.assertAllClose(x, xprime, rtol=0., atol=1.e-4)
    self.assertAllClose(z, zprime, rtol=0., atol=1.e-4)
    self.assertAllClose(
        self.evaluate(-bijection.inverse_log_det_jacobian(x, 3)),
        self.evaluate(bijection.forward_log_det_jacobian(z, 1)),
        rtol=1e-4,
        atol=0.)

  def testBijectiveAndFinite(self):
    """Test that the bijector is invertible and that outputs are finite."""
    bijection = self._create_glow()
    self.evaluate([v.initializer for v in bijection.variables])
    ims = np.zeros([10]+self.output_shape)
    ims += np.linspace(self.minval+0.0001,
                       self.maxval-0.0001,
                       10, dtype=np.float32).reshape([10, 1, 1, 1])

    zs = np.float32(np.zeros([10, np.prod(self.output_shape)]))
    zs += np.linspace(0., 10., 10,
                      dtype=np.float32).reshape(10, 1)  # 10-sigma is far away.

    bijector_test_util.assert_bijective_and_finite(
        bijection, np.float32(zs), np.float32(ims), event_ndims=1,
        eval_func=self.evaluate, inverse_event_ndims=3,
        rtol=1e-3)

  def testDataInit_inverse(self):
    """Test that actnorm data dependent initialization works on inverse pass."""
    bijection = self._create_glow(actnorm=True)
    self.evaluate([v.initializer for v in bijection.variables])
    x = self._make_images()
    nblocks = 0
    made_a_check = False
    splits = bijection.bijectors[1].blockwise_splits
    splits = [[bs[0]+bs[1], bs[2]] for bs in splits]

    x = bijection.bijectors[0].inverse(x)
    for b in bijection.bijectors[1].bijectors:
      if isinstance(b, blockwise._Blockwise):
        x1, x2 = tf.split(x, splits[-2-nblocks], axis=-1)

        for bb in b.bijectors[0].bijectors:
          if isinstance(bb, tfb.glow.ActivationNormalization):
            x1 = self.evaluate(bb.inverse(x1))
            mean = self.evaluate(tf.reduce_mean(x1, axis=(-4, -3, -2)))
            stddev = self.evaluate(tf.math.reduce_std(x1, axis=(-4, -3, -2)))
            self.assertAllClose(mean, tf.zeros_like(mean), atol=1e-5)
            self.assertAllClose(stddev, tf.ones_like(stddev), atol=0, rtol=5e-5)
            made_a_check = True
          else:
            x1 = self.evaluate(bb.inverse(x1))
        x = tf.concat([x1, x2], axis=-1)
        nblocks += 1

      elif isinstance(b, composition.Composition):
        for bb in b.bijectors:
          x = self.evaluate(bb.inverse(x))
          if isinstance(bb, tfb.glow.ActivationNormalization):
            mean = tf.reduce_mean(x, axis=(-4, -3, -2))
            stddev = tf.math.reduce_std(x, axis=(-4, -3, -2))
            self.assertAllClose(mean, tf.zeros_like(mean), atol=1e-5)
            self.assertAllClose(stddev, tf.ones_like(stddev), atol=0, rtol=5e-5)
            made_a_check = True
      else:
        x = self.evaluate(b.inverse(x))

    assert made_a_check

  def testDataInit_forward(self):
    """Test that actnorm data dependent initialization works on forward pass."""
    bijection = self._create_glow(actnorm=True)
    self.evaluate([v.initializer for v in bijection.variables])
    y = np.float32(np.random.normal(0., 1., [5, 768]))
    nblocks = 0
    made_a_check = False
    splits = bijection.bijectors[1].blockwise_splits
    splits = [[bs[0]+bs[1], bs[2]] for bs in splits]

    for b in reversed(bijection.bijectors[1].bijectors):
      if isinstance(b, blockwise._Blockwise):
        y1, y2 = tf.split(y, splits[nblocks], axis=-1)

        for bb in reversed(b.bijectors[0].bijectors):
          if isinstance(bb, tfb.glow.ActivationNormalization):
            y1 = self.evaluate(bb.forward(y1))
            mean = self.evaluate(tf.reduce_mean(y1, axis=(-4, -3, -2)))
            stddev = self.evaluate(tf.math.reduce_std(y1, axis=(-4, -3, -2)))
            self.assertAllClose(mean, tf.zeros_like(mean), atol=1e-5)
            self.assertAllClose(stddev, tf.ones_like(stddev), atol=0, rtol=5e-5)
            made_a_check = True
          else:
            y1 = self.evaluate(bb.forward(y1))
        y = tf.concat([y1, y2], axis=-1)
        nblocks += 1

      elif isinstance(b, composition.Composition):
        for bb in reversed(b.bijectors):
          y = self.evaluate(bb.forward(y))
          if isinstance(bb, tfb.glow.ActivationNormalization):
            mean = tf.reduce_mean(y, axis=(-4, -3, -2))
            stddev = tf.math.reduce_std(y, axis=(-4, -3, -2))
            self.assertAllClose(mean, tf.zeros_like(mean), atol=1e-5)
            self.assertAllClose(stddev, tf.ones_like(stddev), atol=0, rtol=5e-5)
            made_a_check = True
      else:
        y = self.evaluate(b.forward(y))
    assert made_a_check

  def testInverseLogDetJacobian(self):
    """Test if log-det-jacobian agrees with numerical computation."""
    if not tf.executing_eagerly():
      self.skipTest('Theres a problem with numerical computation of Jacobian.'
                    'the bijector jacobian implementation  still returns'
                    'roughly the same values as it does in eager mode, so I'
                    'think our computation works here.')

    x = self._make_images()
    x = tf.constant(x, tf.float32)

    bijection = self._create_glow(actnorm=True)

    self.evaluate([v.initializer for v in bijection.variables])
    jacob_manual = self.evaluate(batch_jacobian(bijection.inverse, x))
    _, ldj_manual = np.linalg.slogdet(jacob_manual.reshape([5, 768, 768]))

    jacob = self.evaluate(bijection.inverse_log_det_jacobian(x, 3))
    self.assertAllClose(ldj_manual, jacob, rtol=1.e-5)

  def testBadBijectorFn(self):
    """Test a bad bijector function.

    Ensures that an error happens if the bijector outputs a tensor whose shape
    does not give a straightforward way to produce a bijector.
    """
    def bad_glow_net(input_shape):
      del input_shape
      return tfb.Pad()

    x = self._make_images()
    x = tf.constant(x, tf.float32)

    bijection = tfb.Glow(
        output_shape=self.output_shape,
        num_glow_blocks=2,
        num_steps_per_block=1,
        coupling_bijector_fn=bad_glow_net,
        exit_bijector_fn=tfb.GlowDefaultExitNetwork,
        grab_after_block=[0.5, 0.5],
        use_actnorm=False)
    self.evaluate([v.initializer for v in bijection.variables])
    try:
      z = bijection.inverse(x)
      del z
      raise Exception('If youre reading this, then something is wrong.')
    except ValueError:
      pass

  def testOtherBadBijectorFn(self):
    """Test another bad bijector function.

    This tests that an error is thrown if the bijector function produces
    something other than a bijector or a tensor
    """
    def other_bad_glow_net(input_shape):
      del input_shape
      return lambda x: tfd.Normal(x, 0)

    x = self._make_images()
    x = tf.constant(x, tf.float32)

    bijection = tfb.Glow(
        output_shape=self.output_shape,
        num_glow_blocks=2,
        num_steps_per_block=1,
        coupling_bijector_fn=other_bad_glow_net,
        exit_bijector_fn=tfb.GlowDefaultExitNetwork,
        grab_after_block=[0.5, 0.5],
        use_actnorm=False)
    self.evaluate([v.initializer for v in bijection.variables])
    try:
      z = bijection.inverse(x)
      del z
      raise Exception('If youre reading this, then something is wrong.')

    except ValueError:
      pass

  def testMultiDimensionalSampling(self):
    """This tests that the model can take samples in variable dimensions."""
    def dummy_glow_net(input_shape):
      del input_shape
      return lambda x: tfb.Identity()

    def dummy_exit_net(input_shape, output_chan):
      del input_shape, output_chan
      return lambda x: tfb.Identity()

    bijection = tfb.Glow(
        output_shape=self.output_shape,
        num_glow_blocks=2,
        num_steps_per_block=1,
        coupling_bijector_fn=dummy_glow_net,
        exit_bijector_fn=dummy_exit_net,
        grab_after_block=[0.5, 0.5],
        use_actnorm=False)
    dist = bijection(tfd.Independent(
        tfd.Normal(tf.zeros([16*16*3]), tf.ones(16*16*3)),
        reinterpreted_batch_ndims=1))

    single_samp2 = dist.sample(1)
    single_samp3 = dist.sample([1, 1])
    single_samp = dist.sample()
    self.assertShapeEqual(np.zeros([16, 16, 3]), single_samp)
    self.assertShapeEqual(np.zeros([1, 16, 16, 3]), single_samp2)
    self.assertShapeEqual(np.zeros([1, 1, 16, 16, 3]), single_samp3)

  def testMultiDimensionalInput(self):
    """This tests if the model runs with different batch shapes."""
    def dummy_glow_net(input_shape):
      del input_shape
      return lambda x: tfb.Identity()

    def dummy_exit_net(input_shape, output_chan):
      del input_shape, output_chan
      return lambda x: tfb.Identity()

    bijection = tfb.Glow(
        output_shape=self.output_shape,
        num_glow_blocks=2,
        num_steps_per_block=1,
        coupling_bijector_fn=dummy_glow_net,
        exit_bijector_fn=dummy_exit_net,
        grab_after_block=[0.5, 0.5],
        use_actnorm=False)

    single_z = bijection.inverse(tf.zeros([16, 16, 3]))
    single_z2 = bijection.inverse(tf.zeros([1, 16, 16, 3]))
    single_z3 = bijection.inverse(tf.zeros([1, 1, 16, 16, 3]))
    self.assertShapeEqual(np.zeros([16 * 16 * 3]), single_z)
    self.assertShapeEqual(np.zeros([1, 16 * 16 * 3]), single_z2)
    self.assertShapeEqual(np.zeros([1, 1, 16 * 16 * 3]), single_z3)

  def testDtypes(self):
    """Test if the bijector identifies dtype correctly."""
    ims = self._make_images()
    bijection = self._create_glow(False)
    z = bijection.inverse(ims)

    def float64_net(input_shape):
      input_nchan = input_shape[-1]
      return tf.keras.Sequential([
          tf.keras.layers.Input(input_shape, dtype=tf.float64),
          tf.keras.layers.Conv2D(
              2 * input_nchan, 3, padding='same', dtype=tf.float64)])
    def float64_exit(input_shape, output_chan):
      return tf.keras.Sequential([
          tf.keras.layers.Input(input_shape, dtype=tf.float64),
          tf.keras.layers.Conv2D(
              2*output_chan, 3, padding='same', dtype=tf.float64)])

    float64_bijection = tfb.Glow(
        output_shape=self.output_shape,
        num_glow_blocks=2,
        num_steps_per_block=1,
        coupling_bijector_fn=float64_net,
        exit_bijector_fn=float64_exit,
        grab_after_block=[0.5, 0.5]
        )
    zf64 = float64_bijection.inverse(tf.cast(ims, tf.float64))
    self.evaluate([v.initializer for v in bijection.variables])
    self.evaluate([v.initializer for v in float64_bijection.variables])
    self.assertAllFinite(self.evaluate(z))
    self.assertAllFinite(self.evaluate(zf64))

  def testBijectorFn(self):
    """Test if the bijector function works for additive coupling."""
    ims = self._make_images()
    def shiftfn(input_shape):
      input_nchan = input_shape[-1]
      return tf.keras.Sequential([
          tf.keras.layers.Input(input_shape),
          tf.keras.layers.Conv2D(
              input_nchan, 3, padding='same')])

    def shiftexitfn(input_shape, output_chan):
      return tf.keras.Sequential([
          tf.keras.layers.Input(input_shape),
          tf.keras.layers.Conv2D(
              output_chan, 3, padding='same')])

    shiftonlyglow = tfb.Glow(
        output_shape=self.output_shape,
        num_glow_blocks=2,
        num_steps_per_block=1,
        coupling_bijector_fn=shiftfn,
        exit_bijector_fn=shiftexitfn,
        grab_after_block=[0.5, 0.5]
        )
    z = shiftonlyglow.inverse(ims)
    self.evaluate([v.initializer for v in shiftonlyglow.variables])
    self.assertAllFinite(self.evaluate(z))


if __name__ == '__main__':
  test_util.main()

# Copyright 2020 The TensorFlow Probability Authors. # # 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 tensorflow_probability.python.bijectors.glow. # Dependency imports Verify that the bijector is properly invertible. # Absolute errors are < 1e-4, and need to use absolute since dealing with # values surrounding 0. Test that the bijector is invertible and that outputs are finite. # 10-sigma is far away. Test that actnorm data dependent initialization works on inverse pass. Test that actnorm data dependent initialization works on forward pass. Test if log-det-jacobian agrees with numerical computation. Test a bad bijector function.      Ensures that an error happens if the bijector outputs a tensor whose shape     does not give a straightforward way to produce a bijector. Test another bad bijector function.      This tests that an error is thrown if the bijector function produces     something other than a bijector or a tensor This tests that the model can take samples in variable dimensions. This tests if the model runs with different batch shapes. Test if the bijector identifies dtype correctly. Test if the bijector function works for additive coupling.