reshape_test.py

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

import numpy as np
import tensorflow.compat.v1 as tf1
import tensorflow.compat.v2 as tf
from tensorflow_probability.python.bijectors import bijector_test_util
from tensorflow_probability.python.bijectors import chain
from tensorflow_probability.python.bijectors import identity
from tensorflow_probability.python.bijectors import reshape
from tensorflow_probability.python.internal import hypothesis_testlib as tfp_hps
from tensorflow_probability.python.internal import test_util


class _ReshapeBijectorTest(object):
  """Base class for testing the reshape transformation.

  Methods defined in this class call a method self.build_shapes() that
  is implemented by subclasses defined below, returning respectively
   ReshapeBijectorTestStatic: static shapes,
   ReshapeBijectorTestDynamic: shape placeholders of known ndims, and
  so that each test in this base class is automatically run over all
  three cases. The subclasses also implement assertRaisesError to test
  for either Python exceptions (in the case of static shapes) or
  TensorFlow op errors (dynamic shapes).
  """

  def testBijector(self):
    """Do a basic sanity check of forward, inverse, jacobian."""
    expected_x = np.random.randn(4, 3, 2)
    expected_y = np.reshape(expected_x, [4, 6])

    shape_in, shape_out = self.build_shapes([3, 2], [6,])
    bijector = reshape.Reshape(
        event_shape_out=shape_out, event_shape_in=shape_in, validate_args=True)
    [
        x_,
        y_,
        fldj_,
        ildj_,
        fest_,
        iest_,
    ] = self.evaluate([
        bijector.inverse(expected_y),
        bijector.forward(expected_x),
        bijector.forward_log_det_jacobian(expected_x, event_ndims=2),
        bijector.inverse_log_det_jacobian(expected_y, event_ndims=2),
        bijector.forward_event_shape_tensor(expected_x.shape),
        bijector.inverse_event_shape_tensor(expected_y.shape),
    ])
    self.assertStartsWith(bijector.name, 'reshape')
    self.assertAllClose(expected_y, y_, rtol=1e-6, atol=0)
    self.assertAllClose(expected_x, x_, rtol=1e-6, atol=0)
    self.assertAllClose(0., fldj_, rtol=1e-6, atol=0)
    self.assertAllClose(0., ildj_, rtol=1e-6, atol=0)
    # Test that event_shape_tensors match fwd/inv result shapes.
    self.assertAllEqual(y_.shape, fest_)
    self.assertAllEqual(x_.shape, iest_)

  def testEventShapeTensor(self):
    """Test event_shape_tensor methods when even ndims may be dynamic."""

    shape_in_static = [2, 3]
    shape_out_static = [6,]
    shape_in, shape_out = self.build_shapes(shape_in_static, shape_out_static)
    bijector = reshape.Reshape(
        event_shape_out=shape_out, event_shape_in=shape_in, validate_args=True)

    # using the _tensor methods, we should always get a fully-specified
    # result since these are evaluated at graph runtime.
    (shape_out_,
     shape_in_) = self.evaluate((
         bijector.forward_event_shape_tensor(shape_in),
         bijector.inverse_event_shape_tensor(shape_out),
     ))
    self.assertAllEqual(shape_out_static, shape_out_)
    self.assertAllEqual(shape_in_static, shape_in_)

  def testScalarReshape(self):
    """Test reshaping to and from a scalar shape ()."""

    expected_x = np.random.randn(4, 3, 1)
    expected_y = np.reshape(expected_x, [4, 3])

    expected_x_scalar = np.random.randn(1,)
    expected_y_scalar = expected_x_scalar[0]

    shape_in, shape_out = self.build_shapes([], [1,])
    bijector = reshape.Reshape(
        event_shape_out=shape_in, event_shape_in=shape_out, validate_args=True)
    (x_,
     y_,
     x_scalar_,
     y_scalar_
    ) = self.evaluate((
        bijector.inverse(expected_y),
        bijector.forward(expected_x),
        bijector.inverse(expected_y_scalar),
        bijector.forward(expected_x_scalar),
    ))
    self.assertAllClose(expected_y, y_, rtol=1e-6, atol=0)
    self.assertAllClose(expected_x, x_, rtol=1e-6, atol=0)
    self.assertAllClose(expected_y_scalar, y_scalar_, rtol=1e-6, atol=0)
    self.assertAllClose(expected_x_scalar, x_scalar_, rtol=1e-6, atol=0)

  def testValidButNonMatchingInputOpError(self):
    x = np.random.randn(4, 3, 2)

    shape_in, shape_out = self.build_shapes([2, 3], [1, 6, 1,])
    bijector = reshape.Reshape(
        event_shape_out=shape_out, event_shape_in=shape_in, validate_args=True)

    # Here we pass in a tensor (x) whose shape is compatible with
    # the output shape, so tf.reshape will throw no error, but
    # doesn't match the expected input shape.
    with self.assertRaisesError('Input `event_shape` does not match'):
      self.evaluate(bijector.forward(x))

  def testValidButNonMatchingInputPartiallySpecifiedOpError(self):
    x = np.random.randn(4, 3, 2)

    shape_in, shape_out = self.build_shapes([2, -1], [1, 6, 1,])
    bijector = reshape.Reshape(
        event_shape_out=shape_out, event_shape_in=shape_in, validate_args=True)

    with self.assertRaisesError('Input `event_shape` does not match'):
      self.evaluate(bijector.forward(x))

  # pylint: disable=invalid-name
  def _testInputOutputMismatchOpError(self, expected_error_message):
    x1 = np.random.randn(4, 2, 3)
    x2 = np.random.randn(4, 1, 1, 5)

    shape_in, shape_out = self.build_shapes([2, 3], [1, 1, 5])
    with self.assertRaisesError(expected_error_message):
      bijector = reshape.Reshape(
          event_shape_out=shape_out,
          event_shape_in=shape_in,
          validate_args=True)
      self.evaluate(bijector.forward(x1))
    with self.assertRaisesError(expected_error_message):
      bijector = reshape.Reshape(
          event_shape_out=shape_out,
          event_shape_in=shape_in,
          validate_args=True)
      self.evaluate(bijector.inverse(x2))
  # pylint: enable=invalid-name

  def testOneShapePartiallySpecified(self):
    expected_x = np.random.randn(4, 6)
    expected_y = np.reshape(expected_x, [4, 2, 3])

    # one of input/output shapes is partially specified
    shape_in, shape_out = self.build_shapes([-1,], [2, 3])
    bijector = reshape.Reshape(
        event_shape_out=shape_out, event_shape_in=shape_in, validate_args=True)
    x_, y_, = self.evaluate([
        bijector.inverse(expected_y),
        bijector.forward(expected_x),
    ])
    self.assertAllClose(expected_y, y_, rtol=1e-6, atol=0)
    self.assertAllClose(expected_x, x_, rtol=1e-6, atol=0)

  def testBothShapesPartiallySpecified(self):
    expected_x = np.random.randn(4, 2, 3)
    expected_y = np.reshape(expected_x, [4, 3, 2])
    shape_in, shape_out = self.build_shapes([-1, 3], [-1, 2])
    bijector = reshape.Reshape(
        event_shape_out=shape_out, event_shape_in=shape_in, validate_args=True)
    x_, y_, = self.evaluate([
        bijector.inverse(expected_y),
        bijector.forward(expected_x),
    ])
    self.assertAllClose(expected_y, y_, rtol=1e-6, atol=0)
    self.assertAllClose(expected_x, x_, rtol=1e-6, atol=0)

  def testDefaultVectorShape(self):
    expected_x = np.random.randn(4, 4)
    expected_y = np.reshape(expected_x, [4, 2, 2])
    _, shape_out = self.build_shapes([-1,], [-1, 2])
    bijector = reshape.Reshape(shape_out, validate_args=True)
    x_, y_, = self.evaluate([
        bijector.inverse(expected_y),
        bijector.forward(expected_x),
    ])
    self.assertAllClose(expected_y, y_, rtol=1e-6, atol=0)
    self.assertAllClose(expected_x, x_, rtol=1e-6, atol=0)

  def build_shapes(self, *args, **kwargs):
    raise NotImplementedError('Subclass failed to implement `build_shapes`.')


@test_util.test_all_tf_execution_regimes
class ReshapeBijectorTestStatic(test_util.TestCase, _ReshapeBijectorTest):

  def build_shapes(self, shape_in, shape_out):
    return shape_in, shape_out

  def assertRaisesError(self, msg):
    return self.assertRaisesRegex(Exception, msg)

  def testEventShape(self):
    shape_in_static = tf.TensorShape([2, 3])
    shape_out_static = tf.TensorShape([6])
    bijector = reshape.Reshape(
        event_shape_out=shape_out_static,
        event_shape_in=shape_in_static,
        validate_args=True)

    # Test that forward_ and inverse_event_shape are correct when
    # event_shape_in/_out are statically known, even when the input shapes
    # are only partially specified.
    self.assertEqual(
        bijector.forward_event_shape(tf.TensorShape([4, 2, 3])).as_list(),
        [4, 6])
    self.assertEqual(
        bijector.inverse_event_shape(tf.TensorShape([4, 6])).as_list(),
        [4, 2, 3])

    # Shape is always known for reshaping in eager mode, so we skip these tests.
    if tf.executing_eagerly():
      return
    self.assertEqual(
        bijector.forward_event_shape(tf.TensorShape([None, 2, 3])).as_list(),
        [None, 6])
    self.assertEqual(
        bijector.inverse_event_shape(tf.TensorShape([None, 6])).as_list(),
        [None, 2, 3])
    # If the input shape is totally unknown, there's nothing we can do!
    self.assertIsNone(
        bijector.forward_event_shape(tf.TensorShape(None)).ndims)

  def testBijectiveAndFinite(self):
    x = np.random.randn(4, 2, 3)
    y = np.reshape(x, [4, 1, 2, 3])
    bijector = reshape.Reshape(
        event_shape_in=[2, 3], event_shape_out=[1, 2, 3], validate_args=True)
    bijector_test_util.assert_bijective_and_finite(
        bijector,
        x,
        y,
        eval_func=self.evaluate,
        event_ndims=2,
        inverse_event_ndims=3,
        rtol=1e-6,
        atol=0)

  def testWorksWithChain(self):
    shape_out = (4,)
    shape_in = (2, 2)
    x = np.zeros(shape_in)
    y = np.zeros(shape_out)
    bijector = chain.Chain([
        identity.Identity(),
        reshape.Reshape(event_shape_out=shape_out, event_shape_in=shape_in)
    ])
    new_y = self.evaluate(bijector.forward(x))
    new_x = self.evaluate(bijector.inverse(y))
    fldj = self.evaluate(
        bijector.forward_log_det_jacobian(x, event_ndims=len(shape_in)))
    ildj = self.evaluate(
        bijector.inverse_log_det_jacobian(y, event_ndims=len(shape_out)))
    self.assertEqual(shape_out, new_y.shape)
    self.assertEqual(shape_in, new_x.shape)
    self.assertEqual((), fldj.shape)
    self.assertEqual((), ildj.shape)

  def testMultipleUnspecifiedDimensionsOpError(self):
    shape_in, shape_out = self.build_shapes([2, 3], [4, -1, -1,])

    with self.assertRaises(ValueError):
      reshape.Reshape(
          event_shape_out=shape_out,
          event_shape_in=shape_in,
          validate_args=True)

  def testInvalidDimensionsOpError(self):
    shape_in, shape_out = self.build_shapes([2, 3], [1, 2, -2,])

    with self.assertRaises(ValueError):
      reshape.Reshape(
          event_shape_out=shape_out,
          event_shape_in=shape_in,
          validate_args=True)

  def testInputOutputMismatchOpError(self):
    self._testInputOutputMismatchOpError('shape')

  def testCheckingVariableShape(self):
    shape_out = tf.Variable([-2, 10])
    self.evaluate(shape_out.initializer)
    with self.assertRaisesOpError(
        'elements must be either positive integers or `-1`'):
      self.evaluate(reshape.Reshape(shape_out, validate_args=True).forward([0]))

  def testCheckingMutatedVariableShape(self):
    shape_out = tf.Variable([1, 1])
    self.evaluate(shape_out.initializer)
    bijector = reshape.Reshape(shape_out, validate_args=True)
    self.evaluate(bijector.forward([0]))
    with self.assertRaisesOpError(
        'elements must be either positive integers or `-1`'):
      with tf.control_dependencies([shape_out.assign([-2, 10])]):
        self.evaluate(bijector.forward([0]))

  @test_util.numpy_disable_test_missing_functionality('b/142265598')
  @test_util.jax_disable_test_missing_functionality('tf.boolean_mask')
  def testConcretizationLimits(self):
    shape_out = tfp_hps.defer_and_count_usage(tf.Variable([1]))
    bijector = reshape.Reshape(shape_out, validate_args=True)
    x = [1]  # Pun: valid input or output, and valid input or output shape
    for method in ['forward', 'inverse', 'forward_event_shape',
                   'inverse_event_shape', 'forward_event_shape_tensor',
                   'inverse_event_shape_tensor']:
      with tfp_hps.assert_no_excessive_var_usage(method, max_permissible=7):
        getattr(bijector, method)(x)
    for method in ['forward_log_det_jacobian', 'inverse_log_det_jacobian']:
      with tfp_hps.assert_no_excessive_var_usage(method, max_permissible=4):
        getattr(bijector, method)(x, event_ndims=1)


class ReshapeBijectorTestDynamic(test_util.TestCase, _ReshapeBijectorTest):

  def build_shapes(self, shape_in, shape_out):
    shape_in = np.array(shape_in, np.int32)
    shape_out = np.array(shape_out, np.int32)
    return (
        tf1.placeholder_with_default(
            shape_in, shape=[len(shape_in)]),
        tf1.placeholder_with_default(
            shape_out, shape=[len(shape_out)]),
    )

  def assertRaisesError(self, msg):
    if tf.executing_eagerly():
      return self.assertRaisesRegex(Exception, msg)
    return self.assertRaisesOpError(msg)

  def testEventShape(self):
    # Shape is always known for reshaping in eager mode, so we skip these tests.
    if tf.executing_eagerly(): return

    event_shape_in, event_shape_out = self.build_shapes([2, 3], [6])
    bijector = reshape.Reshape(
        event_shape_out=event_shape_out,
        event_shape_in=event_shape_in,
        validate_args=True)

    self.assertEqual(
        bijector.forward_event_shape(tf.TensorShape([4, 2, 3])).as_list(),
        [4, None])
    self.assertEqual(
        bijector.forward_event_shape(tf.TensorShape([None, 2, 3])).as_list(),
        [None, None])
    self.assertEqual(
        bijector.inverse_event_shape(tf.TensorShape([4, 6])).as_list(),
        [4, None, None])
    self.assertEqual(
        bijector.inverse_event_shape(tf.TensorShape([None, 6])).as_list(),
        [None, None, None])
    # If the input shape is totally unknown, there's nothing we can do!
    self.assertIsNone(
        bijector.forward_event_shape(tf.TensorShape(None)).ndims)

  def testInputOutputMismatchOpError(self):
    self._testInputOutputMismatchOpError('shape')

  def testMultipleUnspecifiedDimensionsOpError(self):
    with self.assertRaisesError('must have at most one `-1`.'):
      shape_in, shape_out = self.build_shapes([2, 3], [4, -1, -1,])
      bijector = reshape.Reshape(
          event_shape_out=shape_out,
          event_shape_in=shape_in,
          validate_args=True)
      self.evaluate(bijector.forward_event_shape_tensor(shape_in))

  def testInvalidDimensionsOpError(self):
    shape_in, shape_out = self.build_shapes([2, 3], [1, 2, -2,])

    with self.assertRaisesError(
        'elements must be either positive integers or `-1`.'):
      bijector = reshape.Reshape(
          event_shape_out=shape_out,
          event_shape_in=shape_in,
          validate_args=True)
      self.evaluate(bijector.forward_event_shape_tensor(shape_in))

  def testUnknownShapeRank(self):
    if tf.executing_eagerly(): return
    unknown_shape = tf1.placeholder_with_default([2, 2], shape=None)
    known_shape = [2, 2]

    with self.assertRaisesRegex(NotImplementedError,
                                'must be statically known.'):
      reshape.Reshape(event_shape_out=unknown_shape)

    with self.assertRaisesRegex(NotImplementedError,
                                'must be statically known.'):
      reshape.Reshape(event_shape_out=known_shape, event_shape_in=unknown_shape)

  def testScalarInVectorOut(self):
    bijector = reshape.Reshape(event_shape_in=[], event_shape_out=[-1])
    self.assertAllEqual(np.zeros([3, 4, 5, 1]),
                        self.evaluate(bijector.forward(np.zeros([3, 4, 5]))))
    self.assertAllEqual(np.zeros([3, 4, 5]),
                        self.evaluate(bijector.inverse(np.zeros([3, 4, 5, 1]))))


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