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!")



[((18, 18, 18, 50), 'neural_spline_flows.nde.transforms.base.InverseTransform', 'base.InverseTransform', ({(18, 40, 18, 49): 'transform'}, {}), '(transform)', False, 'from neural_spline_flows.nde.transforms import base\n'), ((19, 19, 19, 64), 'neural_spline_flows.nde.transforms.base.CompositeTransform', 'base.CompositeTransform', ({(19, 43, 19, 63): '[inverse, transform]'}, {}), '([inverse, transform])', False, 'from neural_spline_flows.nde.transforms import base\n'), ((25, 36, 25, 65), 'torch.zeros', 'torch.zeros', ({(25, 48, 25, 64): 'inputs.shape[:1]'}, {}), '(inputs.shape[:1])', False, 'import torch\n'), ((15, 43, 15, 60), 'torch.Size', 'torch.Size', ({(15, 54, 15, 59): 'shape'}, {}), '(shape)', False, 'import torch\n'), ((29, 35, 29, 61), 'torch.equal', 'torch.equal', ({(29, 47, 29, 52): 'first', (29, 54, 29, 60): 'second'}, {}), '(first, second)', False, 'import torch\n'), ((12, 25, 12, 44), 'torch.isnan', 'torch.isnan', ({(12, 37, 12, 43): 'tensor'}, {}), '(tensor)', False, 'import torch\n'), ((13, 25, 13, 44), 'torch.isinf', 'torch.isinf', ({(13, 37, 13, 43): 'tensor'}, {}), '(tensor)', False, 'import torch\n')]