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, 'neural_spline_flows.nde.transforms.base.InverseTransform', 'base.InverseTransform', (['transform'], {}), False, 'from neural_spline_flows.nde.transforms import base\n'), (19, 'neural_spline_flows.nde.transforms.base.CompositeTransform', 'base.CompositeTransform', (['[inverse, transform]'], {}), False, 'from neural_spline_flows.nde.transforms import base\n'), (25, 'torch.zeros', 'torch.zeros', (['inputs.shape[:1]'], {}), False, 'import torch\n'), (15, 'torch.Size', 'torch.Size', (['shape'], {}), False, 'import torch\n'), (29, 'torch.equal', 'torch.equal', (['first', 'second'], {}), False, 'import torch\n'), (12, 'torch.isnan', 'torch.isnan', (['tensor'], {}), False, 'import torch\n'), (13, 'torch.isinf', 'torch.isinf', (['tensor'], {}), False, 'import torch\n')]