import torch


def is_invertible_module(module_in, test_input_shape, test_input_dtype=torch.float32, atol=1e-6):
    """Test if a :obj:`torch.nn.Module` is invertible

    Parameters
    ----------
    module_in : :obj:`torch.nn.Module`
        A torch.nn.Module to test.
    test_input_shape : :obj:`tuple`
        Dimensions of test tensor object to perform the test with.
    test_input_dtype : :obj:`torch.dtype`, optional
        Data type of test tensor object to perform the test with.
    atol : :obj:`float`, optional
        Tolerance value used for comparing the outputs

    Returns
    -------
        :obj:`bool`
            True if the input module is invertible, False otherwise.

    """
    test_input = torch.rand(test_input_shape, dtype=test_input_dtype)
    if not hasattr(module_in, "inverse"):
        return False
    with torch.no_grad():
        if not torch.allclose(module_in.inverse(module_in(test_input)), test_input, atol=atol):
            return False
        if not torch.allclose(module_in(module_in.inverse(test_input)), test_input, atol=atol):
            return False
        if test_input is module_in(test_input):
            return False
        if test_input is module_in.inverse(test_input):
            return False
    return True