# Copyright 2018 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 Bijector."""

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

import tensorflow.compat.v2 as tf
from tensorflow_probability.python import bijectors as tfb
from tensorflow_probability.python import distributions as tfd

from tensorflow_probability.python.bijectors import bijector_test_util
from tensorflow_probability.python.internal import tensorshape_util
from tensorflow_probability.python.internal import test_util as tfp_test_util
from tensorflow.python.framework import test_util  # pylint: disable=g-direct-tensorflow-import


@test_util.run_all_in_graph_and_eager_modes
class InvertBijectorTest(tf.test.TestCase):
  """Tests the correctness of the Y = Invert(bij) transformation."""

  def testBijector(self):
    for fwd in [
        tfb.Identity(),
        tfb.Exp(),
        tfb.Affine(shift=[0., 1.], scale_diag=[2., 3.]),
        tfb.Softplus(),
        tfb.SoftmaxCentered(),
    ]:
      rev = tfb.Invert(fwd)
      self.assertStartsWith(rev.name, "_".join(["invert", fwd.name]))
      x = [[[1., 2.],
            [2., 3.]]]
      self.assertAllClose(
          self.evaluate(fwd.inverse(x)), self.evaluate(rev.forward(x)))
      self.assertAllClose(
          self.evaluate(fwd.forward(x)), self.evaluate(rev.inverse(x)))
      self.assertAllClose(
          self.evaluate(fwd.forward_log_det_jacobian(x, event_ndims=1)),
          self.evaluate(rev.inverse_log_det_jacobian(x, event_ndims=1)))
      self.assertAllClose(
          self.evaluate(fwd.inverse_log_det_jacobian(x, event_ndims=1)),
          self.evaluate(rev.forward_log_det_jacobian(x, event_ndims=1)))

  def testScalarCongruency(self):
    bijector = tfb.Invert(tfb.Exp())
    bijector_test_util.assert_scalar_congruency(
        bijector, lower_x=1e-3, upper_x=1.5, eval_func=self.evaluate, rtol=0.05)

  def testShapeGetters(self):
    bijector = tfb.Invert(
        tfb.SoftmaxCentered(validate_args=True))
    x = tf.TensorShape([2])
    y = tf.TensorShape([1])
    self.assertAllEqual(y, bijector.forward_event_shape(x))
    self.assertAllEqual(
        tensorshape_util.as_list(y),
        self.evaluate(
            bijector.forward_event_shape_tensor(tensorshape_util.as_list(x))))
    self.assertAllEqual(x, bijector.inverse_event_shape(y))
    self.assertAllEqual(
        tensorshape_util.as_list(x),
        self.evaluate(
            bijector.inverse_event_shape_tensor(tensorshape_util.as_list(y))))

  def testDocstringExample(self):
    exp_gamma_distribution = (
        tfd.TransformedDistribution(
            distribution=tfd.Gamma(concentration=1., rate=2.),
            bijector=tfb.Invert(tfb.Exp())))
    self.assertAllEqual(
        [],
        self.evaluate(
            tf.shape(
                exp_gamma_distribution.sample(seed=tfp_test_util.test_seed()))))


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

[((92, 2, 92, 16), 'tensorflow.compat.v2.test.main', 'tf.test.main', ({}, {}), '()', True, 'import tensorflow.compat.v2 as tf\n'), ((60, 4, 61, 80), 'tensorflow_probability.python.bijectors.bijector_test_util.assert_scalar_congruency', 'bijector_test_util.assert_scalar_congruency', (), '', False, 'from tensorflow_probability.python.bijectors import bijector_test_util\n'), ((66, 8, 66, 27), 'tensorflow.compat.v2.TensorShape', 'tf.TensorShape', ({(66, 23, 66, 26): '[2]'}, {}), '([2])', True, 'import tensorflow.compat.v2 as tf\n'), ((67, 8, 67, 27), 'tensorflow.compat.v2.TensorShape', 'tf.TensorShape', ({(67, 23, 67, 26): '[1]'}, {}), '([1])', True, 'import tensorflow.compat.v2 as tf\n'), ((37, 8, 37, 22), 'tensorflow_probability.python.bijectors.Identity', 'tfb.Identity', ({}, {}), '()', True, 'from tensorflow_probability.python import bijectors as tfb\n'), ((38, 8, 38, 17), 'tensorflow_probability.python.bijectors.Exp', 'tfb.Exp', ({}, {}), '()', True, 'from tensorflow_probability.python import bijectors as tfb\n'), ((39, 8, 39, 55), 'tensorflow_probability.python.bijectors.Affine', 'tfb.Affine', (), '', True, 'from tensorflow_probability.python import bijectors as tfb\n'), ((40, 8, 40, 22), 'tensorflow_probability.python.bijectors.Softplus', 'tfb.Softplus', ({}, {}), '()', True, 'from tensorflow_probability.python import bijectors as tfb\n'), ((41, 8, 41, 29), 'tensorflow_probability.python.bijectors.SoftmaxCentered', 'tfb.SoftmaxCentered', ({}, {}), '()', True, 'from tensorflow_probability.python import bijectors as tfb\n'), ((43, 12, 43, 27), 'tensorflow_probability.python.bijectors.Invert', 'tfb.Invert', ({(43, 23, 43, 26): 'fwd'}, {}), '(fwd)', True, 'from tensorflow_probability.python import bijectors as tfb\n'), ((59, 26, 59, 35), 'tensorflow_probability.python.bijectors.Exp', 'tfb.Exp', ({}, {}), '()', True, 'from tensorflow_probability.python import bijectors as tfb\n'), ((65, 8, 65, 47), 'tensorflow_probability.python.bijectors.SoftmaxCentered', 'tfb.SoftmaxCentered', (), '', True, 'from tensorflow_probability.python import bijectors as tfb\n'), ((70, 8, 70, 35), 'tensorflow_probability.python.internal.tensorshape_util.as_list', 'tensorshape_util.as_list', ({(70, 33, 70, 34): 'y'}, {}), '(y)', False, 'from tensorflow_probability.python.internal import tensorshape_util\n'), ((75, 8, 75, 35), 'tensorflow_probability.python.internal.tensorshape_util.as_list', 'tensorshape_util.as_list', ({(75, 33, 75, 34): 'x'}, {}), '(x)', False, 'from tensorflow_probability.python.internal import tensorshape_util\n'), ((82, 25, 82, 61), 'tensorflow_probability.python.distributions.Gamma', 'tfd.Gamma', (), '', True, 'from tensorflow_probability.python import distributions as tfd\n'), ((72, 48, 72, 75), 'tensorflow_probability.python.internal.tensorshape_util.as_list', 'tensorshape_util.as_list', ({(72, 73, 72, 74): 'x'}, {}), '(x)', False, 'from tensorflow_probability.python.internal import tensorshape_util\n'), ((77, 48, 77, 75), 'tensorflow_probability.python.internal.tensorshape_util.as_list', 'tensorshape_util.as_list', ({(77, 73, 77, 74): 'y'}, {}), '(y)', False, 'from tensorflow_probability.python.internal import tensorshape_util\n'), ((83, 32, 83, 41), 'tensorflow_probability.python.bijectors.Exp', 'tfb.Exp', ({}, {}), '()', True, 'from tensorflow_probability.python import bijectors as tfb\n'), ((88, 51, 88, 76), 'tensorflow_probability.python.internal.test_util.test_seed', 'tfp_test_util.test_seed', ({}, {}), '()', True, 'from tensorflow_probability.python.internal import test_util as tfp_test_util\n')]