import torch
import torchtestcase

from neural_spline_flows.nde.transforms import base


class TransformTest(torchtestcase.TorchTestCase):
    """Base test for all transforms."""

    def assert_tensor_is_good(self, tensor, shape=None):
        self.assertIsInstance(tensor, torch.Tensor)
        self.assertFalse(torch.isnan(tensor).any())
        self.assertFalse(torch.isinf(tensor).any())
        if shape is not None:
            self.assertEqual(tensor.shape, torch.Size(shape))

    def assert_forward_inverse_are_consistent(self, transform, inputs):
        inverse = base.InverseTransform(transform)
        identity = base.CompositeTransform([inverse, transform])
        outputs, logabsdet = identity(inputs)

        self.assert_tensor_is_good(outputs, shape=inputs.shape)
        self.assert_tensor_is_good(logabsdet, shape=inputs.shape[:1])
        self.assertEqual(outputs, inputs)
        self.assertEqual(logabsdet, torch.zeros(inputs.shape[:1]))

    def assertNotEqual(self, first, second, msg=None):
        if ((self._eps and (first - second).abs().max().item() < self._eps) or
                (not self._eps and torch.equal(first, second))):
            self._fail_with_message(msg, "The tensors are _not_ different!")



[((576, 608), 'neural_spline_flows.nde.transforms.base.InverseTransform', 'base.InverseTransform', (['transform'], {}), '(transform)\n', (597, 608), False, 'from neural_spline_flows.nde.transforms import base\n'), ((628, 673), 'neural_spline_flows.nde.transforms.base.CompositeTransform', 'base.CompositeTransform', (['[inverse, transform]'], {}), '([inverse, transform])\n', (651, 673), False, 'from neural_spline_flows.nde.transforms import base\n'), ((933, 962), 'torch.zeros', 'torch.zeros', (['inputs.shape[:1]'], {}), '(inputs.shape[:1])\n', (944, 962), False, 'import torch\n'), ((466, 483), 'torch.Size', 'torch.Size', (['shape'], {}), '(shape)\n', (476, 483), False, 'import torch\n'), ((1134, 1160), 'torch.equal', 'torch.equal', (['first', 'second'], {}), '(first, second)\n', (1145, 1160), False, 'import torch\n'), ((314, 333), 'torch.isnan', 'torch.isnan', (['tensor'], {}), '(tensor)\n', (325, 333), False, 'import torch\n'), ((366, 385), 'torch.isinf', 'torch.isinf', (['tensor'], {}), '(tensor)\n', (377, 385), False, 'import torch\n')]