deepinv 0.1.0.dev0__py3-none-any.whl

deepinv/tests/test_physics.py

@@ -0,0 +1,316 @@
+import pytest
+import torch
+import numpy as np
+
+import deepinv as dinv
+
+
+# Linear forward operators to test (make sure they appear in find_operator as well)
+OPERATORS = [
+    "CS",
+    "fastCS",
+    "inpainting",
+    "denoising",
+    "deblur_fft",
+    "deblur",
+    "singlepixel",
+    "fast_singlepixel",
+    "super_resolution",
+    "MRI",
+    "pansharpen",
+]
+NONLINEAR_OPERATORS = ["haze", "blind_deblur", "lidar"]
+
+NOISES = [
+    "Gaussian",
+    "Poisson",
+    "PoissonGaussian",
+    "UniformGaussian",
+    "Uniform",
+    "Neighbor2Neighbor",
+]
+
+
+def find_operator(name, device):
+    r"""
+    Chooses operator
+
+    :param name: operator name
+    :param device: (torch.device) cpu or cuda
+    :return: (deepinv.physics.Physics) forward operator.
+    """
+    img_size = (3, 16, 8)
+    norm = 1
+    if name == "CS":
+        m = 30
+        p = dinv.physics.CompressedSensing(m=m, img_shape=img_size, device=device)
+        norm = (
+            1 + np.sqrt(np.prod(img_size) / m)
+        ) ** 2 - 0.75  # Marcenko-Pastur law, second term is a small n correction
+    elif name == "fastCS":
+        p = dinv.physics.CompressedSensing(
+            m=20, fast=True, channelwise=True, img_shape=img_size, device=device
+        )
+    elif name == "inpainting":
+        p = dinv.physics.Inpainting(tensor_size=img_size, mask=0.5, device=device)
+    elif name == "MRI":
+        img_size = (2, 16, 8)
+        p = dinv.physics.MRI(mask=torch.ones(img_size[-2], img_size[-1]), device=device)
+    elif name == "Tomography":
+        img_size = (1, 16, 16)
+        p = dinv.physics.Tomography(
+            img_width=img_size[-1], angles=img_size[-1], device=device
+        )
+    elif name == "denoising":
+        p = dinv.physics.Denoising(dinv.physics.GaussianNoise(0.1))
+    elif name == "pansharpen":
+        img_size = (3, 30, 32)
+        p = dinv.physics.Pansharpen(img_size=img_size, device=device)
+        norm = 0.4
+    elif name == "fast_singlepixel":
+        p = dinv.physics.SinglePixelCamera(
+            m=20, fast=True, img_shape=img_size, device=device
+        )
+    elif name == "singlepixel":
+        m = 20
+        p = dinv.physics.SinglePixelCamera(
+            m=m, fast=False, img_shape=img_size, device=device
+        )
+        norm = (
+            1 + np.sqrt(np.prod(img_size) / m)
+        ) ** 2 - 3.7  # Marcenko-Pastur law, second term is a small n correction
+    elif name == "deblur":
+        img_size = (3, 17, 19)
+        p = dinv.physics.Blur(
+            dinv.physics.blur.gaussian_blur(sigma=(2, 0.1), angle=45.0), device=device
+        )
+    elif name == "deblur_fft":
+        img_size = (3, 17, 19)
+        p = dinv.physics.BlurFFT(
+            img_size=img_size,
+            filter=dinv.physics.blur.gaussian_blur(sigma=(0.1, 0.5), angle=45.0),
+            device=device,
+        )
+    elif name == "super_resolution":
+        img_size = (1, 32, 32)
+        factor = 2
+        norm = 1 / factor**2
+        p = dinv.physics.Downsampling(img_size=img_size, factor=factor, device=device)
+    else:
+        raise Exception("The inverse problem chosen doesn't exist")
+    return p, img_size, norm
+
+
+def find_nonlinear_operator(name, device):
+    r"""
+    Chooses operator
+
+    :param name: operator name
+    :param device: (torch.device) cpu or cuda
+    :return: (deepinv.physics.Physics) forward operator.
+    """
+    if name == "blind_deblur":
+        x = dinv.utils.TensorList(
+            [
+                torch.randn(1, 3, 16, 16, device=device),
+                torch.randn(1, 1, 3, 3, device=device),
+            ]
+        )
+        p = dinv.physics.BlindBlur(kernel_size=3)
+    elif name == "haze":
+        x = dinv.utils.TensorList(
+            [
+                torch.randn(1, 1, 16, 16, device=device),
+                torch.randn(1, 1, 16, 16, device=device),
+                torch.randn(1, device=device),
+            ]
+        )
+        p = dinv.physics.Haze()
+    elif name == "lidar":
+        x = torch.rand(1, 3, 16, 16, device=device)
+        p = dinv.physics.SinglePhotonLidar(device=device)
+    else:
+        raise Exception("The inverse problem chosen doesn't exist")
+    return p, x
+
+
+@pytest.mark.parametrize("name", OPERATORS)
+def test_operators_adjointness(name, device):
+    r"""
+    Tests if a linear forward operator has a well defined adjoint.
+    Warning: Only test linear operators, non-linear ones will fail the test.
+
+    :param name: operator name (see find_operator)
+    :param imsize: image size tuple in (C, H, W)
+    :param device: (torch.device) cpu or cuda:x
+    :return: asserts adjointness
+    """
+    physics, imsize, _ = find_operator(name, device)
+    x = torch.randn(imsize, device=device).unsqueeze(0)
+    error = physics.adjointness_test(x).abs()
+    assert error < 1e-3
+
+
+@pytest.mark.parametrize("name", OPERATORS)
+def test_operators_norm(name, device):
+    r"""
+    Tests if a linear physics operator has a norm close to 1.
+    Warning: Only test linear operators, non-linear ones will fail the test.
+
+    :param name: operator name (see find_operator)
+    :param imsize: (tuple) image size tuple in (C, H, W)
+    :param device: (torch.device) cpu or cuda:x
+    :return: asserts norm is in (.8,1.2)
+    """
+    if name == "singlepixel" or name == "CS":
+        device = torch.device("cpu")
+
+    torch.manual_seed(0)
+    physics, imsize, norm_ref = find_operator(name, device)
+    x = torch.randn(imsize, device=device).unsqueeze(0)
+    norm = physics.compute_norm(x)
+    assert torch.abs(norm - norm_ref) < 0.2
+
+
+@pytest.mark.parametrize("name", NONLINEAR_OPERATORS)
+def test_nonlinear_operators(name, device):
+    r"""
+    Tests if a linear physics operator has a norm close to 1.
+    Warning: Only test linear operators, non-linear ones will fail the test.
+
+    :param name: operator name (see find_operator)
+    :param device: (torch.device) cpu or cuda:x
+    :return: asserts correct shapes
+    """
+    physics, x = find_nonlinear_operator(name, device)
+    y = physics(x)
+    xhat = physics.A_dagger(y)
+    assert x.shape == xhat.shape
+
+
+@pytest.mark.parametrize("name", OPERATORS)
+def test_pseudo_inverse(name, device):
+    r"""
+    Tests if a linear physics operator has a well defined pseudoinverse.
+    Warning: Only test linear operators, non-linear ones will fail the test.
+
+    :param name: operator name (see find_operator)
+    :param imsize: (tuple) image size tuple in (C, H, W)
+    :param device: (torch.device) cpu or cuda:x
+    :return: asserts error is less than 1e-3
+    """
+    physics, imsize, _ = find_operator(name, device)
+    x = torch.randn(imsize, device=device).unsqueeze(0)
+
+    r = physics.A_adjoint(physics.A(x))
+    y = physics.A(r)
+    error = (physics.A_dagger(y) - r).flatten().mean().abs()
+    assert error < 0.01
+
+
+def test_MRI(device):
+    r"""
+    Test MRI function
+
+    :param name: operator name (see find_operator)
+    :param imsize: (tuple) image size tuple in (C, H, W)
+    :param device: (torch.device) cpu or cuda:x
+    :return: asserts error is less than 1e-3
+    """
+    physics = dinv.physics.MRI(mask=None, device=device, acceleration_factor=4)
+    x = torch.randn((2, 320, 320), device=device).unsqueeze(0)
+    x2 = physics.A_adjoint(physics.A(x))
+    assert x2.shape == x.shape
+
+    physics = dinv.physics.MRI(mask=None, device=device, acceleration_factor=8, seed=0)
+    y1 = physics.A(x)
+    physics.reset()
+    y2 = physics.A(x)
+    if y1.shape == y2.shape:
+        error = (y1.abs() - y2.abs()).flatten().mean().abs()
+        assert error > 0.0
+
+
+def choose_noise(noise_type):
+    gain = 0.1
+    sigma = 0.1
+    if noise_type == "PoissonGaussian":
+        noise_model = dinv.physics.PoissonGaussianNoise(sigma=sigma, gain=gain)
+    elif noise_type == "Gaussian":
+        noise_model = dinv.physics.GaussianNoise(sigma)
+    elif noise_type == "UniformGaussian":
+        noise_model = dinv.physics.UniformGaussianNoise(
+            sigma=sigma
+        )  # This is equivalent to GaussianNoise when sigma is fixed
+    elif noise_type == "Uniform":
+        noise_model = dinv.physics.UniformNoise(a=gain)
+    elif noise_type == "Poisson":
+        noise_model = dinv.physics.PoissonNoise(gain)
+    elif noise_type == "Neighbor2Neighbor":
+        noise_model = dinv.physics.PoissonNoise(gain)
+    else:
+        raise Exception("Noise model not found")
+
+    return noise_model
+
+
+@pytest.mark.parametrize("noise_type", NOISES)
+def test_noise(device, noise_type):
+    r"""
+    Tests noise models.
+    """
+    physics = dinv.physics.DecomposablePhysics()
+    physics.noise_model = choose_noise(noise_type)
+    x = torch.ones((1, 12, 7), device=device).unsqueeze(0)
+
+    y1 = physics(
+        x
+    )  # Note: this works but not physics.A(x) because only the noise is reset (A does not encapsulate noise)
+    assert y1.shape == x.shape
+
+    if noise_type == "UniformGaussian":
+        physics.reset()
+        y2 = physics(x)
+        error = (y1 - y2).flatten().abs().sum()
+        assert error > 0.0
+
+
+def test_reset_noise(device):
+    r"""
+    Tests that the reset function works.
+
+    :param device: (torch.device) cpu or cuda:x
+    :return: asserts error is > 0
+    """
+    physics = dinv.physics.DecomposablePhysics()
+    physics.noise_model = dinv.physics.UniformGaussianNoise(
+        sigma=None
+    )  # Should be 20/255 (to check)
+    x = torch.ones((1, 12, 7), device=device).unsqueeze(0)
+
+    y1 = physics(
+        x
+    )  # Note: this works but not physics.A(x) because only the noise is reset (A does not encapsulate noise)
+    physics.reset()
+    y2 = physics(x)
+    error = (y1 - y2).flatten().abs().sum()
+    assert error > 0.0
+
+
+def test_tomography(device):
+    r"""
+    Tests tomography operator which does not have a numerically precise adjoint.
+
+    :param device: (torch.device) cpu or cuda:x
+    """
+    for circle in [True, False]:
+        imsize = (1, 16, 16)
+        physics = dinv.physics.Tomography(
+            img_width=imsize[-1], angles=imsize[-1], device=device, circle=circle
+        )
+
+        x = torch.randn(imsize, device=device).unsqueeze(0)
+        r = physics.A_adjoint(physics.A(x))
+        y = physics.A(r)
+        error = (physics.A_dagger(y) - r).flatten().mean().abs()
+        assert error < 0.2

deepinv/tests/test_sampling.py

@@ -0,0 +1,158 @@
+import pytest
+import torch.nn
+import numpy as np
+
+import deepinv as dinv
+from deepinv.optim.data_fidelity import L2
+from deepinv.sampling import ULA, SKRock, DiffPIR, DPS
+
+
+SAMPLING_ALGOS = ["DDRM", "ULA", "SKRock"]
+
+
+def choose_algo(algo, likelihood, thresh_conv, sigma, sigma_prior):
+    if algo == "ULA":
+        out = ULA(
+            GaussianScore(sigma_prior),
+            likelihood,
+            max_iter=500,
+            thinning=1,
+            step_size=0.01 / (1 / sigma**2 + 1 / sigma_prior**2),
+            clip=(-100, 100),
+            thresh_conv=thresh_conv,
+            sigma=1,
+            verbose=True,
+        )
+    elif algo == "SKRock":
+        out = SKRock(
+            GaussianScore(sigma_prior),
+            likelihood,
+            max_iter=500,
+            inner_iter=5,
+            step_size=1 / (1 / sigma**2 + 1 / sigma_prior**2),
+            clip=(-100, 100),
+            thresh_conv=thresh_conv,
+            sigma=1,
+            verbose=True,
+        )
+    elif algo == "DDRM":
+        diff = dinv.sampling.DDRM(
+            denoiser=GaussianDenoiser(sigma_prior),
+            eta=1,
+            sigmas=np.linspace(1, 0, 100),
+        )
+        out = dinv.sampling.DiffusionSampler(diff, clip=(-100, 100), max_iter=500)
+    else:
+        raise Exception("The sampling algorithm doesnt exist")
+
+    return out
+
+
+class GaussianScore(torch.nn.Module):
+    def __init__(self, sigma_prior):
+        super().__init__()
+        self.sigma_prior2 = sigma_prior**2
+
+    def forward(self, x, sigma):
+        return x / self.sigma_prior2
+
+
+class GaussianDenoiser(torch.nn.Module):
+    def __init__(self, sigma_prior):
+        super().__init__()
+        self.sigma_prior2 = sigma_prior**2
+
+    def forward(self, x, sigma):
+        return x / (1 + sigma**2 / self.sigma_prior2)
+
+
+@pytest.mark.parametrize("algo", SAMPLING_ALGOS)
+def test_sampling_algo(algo, imsize, device):
+    test_sample = torch.ones((1, *imsize))
+
+    sigma = 1
+    sigma_prior = 1
+    physics = dinv.physics.Denoising()
+    physics.noise_model = dinv.physics.GaussianNoise(sigma)
+    y = physics(test_sample)
+
+    convergence_crit = 0.1  # for fast tests
+    likelihood = L2(sigma=sigma)
+    f = choose_algo(
+        algo,
+        likelihood,
+        thresh_conv=convergence_crit,
+        sigma=sigma,
+        sigma_prior=sigma_prior,
+    )
+
+    xmean, xvar = f(y, physics, seed=0)
+
+    tol = 5  # can be lowered?
+    sigma2 = sigma**2
+    sigma_prior2 = sigma_prior**2
+
+    # the posterior of a gaussian likelihood with a gaussian prior is gaussian
+    post_var = (sigma2 * sigma_prior2) / (sigma2 + sigma_prior2)
+    post_mean = y / (1 + sigma2 / sigma_prior2)
+
+    mean_ok = (
+        torch.sum((xmean - post_mean).abs() / post_mean < tol)
+        > np.prod(xmean.shape) / 2
+    )
+
+    var_ok = (
+        torch.sum((xvar - post_var).abs() / post_var < tol) > np.prod(xvar.shape) / 2
+    )
+
+    assert f.mean_has_converged() and f.var_has_converged() and mean_ok and var_ok
+
+
+def test_diffpir(device):
+    from deepinv.models import DiffUNet
+
+    x = torch.ones((1, 3, 32, 32)).to(device)
+
+    sigma = 12.75 / 255.0  # noise level
+
+    physics = dinv.physics.BlurFFT(
+        img_size=(3, x.shape[-2], x.shape[-1]),
+        filter=torch.ones((1, 1, 5, 5), device=device) / 25,
+        device=device,
+        noise_model=dinv.physics.GaussianNoise(sigma=sigma),
+    )
+
+    y = physics(x)
+
+    model = DiffUNet().to(device)
+    likelihood = L2()
+
+    algorithm = DiffPIR(model, likelihood, max_iter=5, verbose=False, device=device)
+
+    out = algorithm(y, physics)
+    assert out.shape == x.shape
+
+
+def test_dps(device):
+    from deepinv.models import DiffUNet
+
+    x = torch.ones((1, 3, 32, 32)).to(device)
+
+    sigma = 12.75 / 255.0  # noise level
+
+    physics = dinv.physics.BlurFFT(
+        img_size=(3, x.shape[-2], x.shape[-1]),
+        filter=torch.ones((1, 1, 5, 5), device=device) / 25,
+        device=device,
+        noise_model=dinv.physics.GaussianNoise(sigma=sigma),
+    )
+
+    y = physics(x)
+
+    model = DiffUNet().to(device)
+    likelihood = L2()
+
+    algorithm = DPS(model, likelihood, max_iter=5, verbose=False, device=device)
+
+    out = algorithm(y, physics)
+    assert out.shape == x.shape

deepinv/tests/test_unfolded.py

@@ -0,0 +1,158 @@
+import pytest
+import torch
+
+import deepinv as dinv
+from deepinv.optim.prior import PnP
+from deepinv.optim.data_fidelity import L2
+from deepinv.unfolded import unfolded_builder, DEQ_builder
+
+
+OPTIM_ALGO = ["PGD", "HQS"]
+
+
+@pytest.mark.parametrize("unfolded_algo", OPTIM_ALGO)
+def test_unfolded(unfolded_algo, imsize, dummy_dataset, device):
+    pytest.importorskip("pytorch_wavelets")
+
+    # Select the data fidelity term
+    data_fidelity = L2()
+
+    # Set up the trainable denoising prior; here, the soft-threshold in a wavelet basis.
+    # If the prior is initialized with a list of length max_iter,
+    # then a distinct weight is trained for each PGD iteration.
+    # For fixed trained model prior across iterations, initialize with a single model.
+    max_iter = 30 if torch.cuda.is_available() else 20  # Number of unrolled iterations
+    level = 3
+    prior = [
+        PnP(denoiser=dinv.models.WaveletPrior(wv="db8", level=level, device=device))
+        for i in range(max_iter)
+    ]
+
+    # Unrolled optimization algorithm parameters
+    lamb = [
+        1.0
+    ] * max_iter  # initialization of the regularization parameter. A distinct lamb is trained for each iteration.
+    stepsize = [
+        1.0
+    ] * max_iter  # initialization of the stepsizes. A distinct stepsize is trained for each iteration.
+
+    sigma_denoiser_init = 0.01
+    sigma_denoiser = [sigma_denoiser_init * torch.ones(level, 3)] * max_iter
+    # sigma_denoiser = [torch.Tensor([sigma_denoiser_init])]*max_iter
+    params_algo = {  # wrap all the restoration parameters in a 'params_algo' dictionary
+        "stepsize": stepsize,
+        "g_param": sigma_denoiser,
+        "lambda": lamb,
+    }
+
+    trainable_params = [
+        "g_param",
+        "stepsize",
+    ]  # define which parameters from 'params_algo' are trainable
+
+    # Define the unfolded trainable model.
+    model = unfolded_builder(
+        unfolded_algo,
+        params_algo=params_algo,
+        trainable_params=trainable_params,
+        data_fidelity=data_fidelity,
+        max_iter=max_iter,
+        prior=prior,
+    )
+
+    for idx, (name, param) in enumerate(model.named_parameters()):
+        assert param.requires_grad
+        assert (trainable_params[0] in name) or (trainable_params[1] in name)
+
+
+@pytest.mark.parametrize("unfolded_algo", OPTIM_ALGO)
+def test_DEQ(unfolded_algo, imsize, dummy_dataset, device):
+    pytest.importorskip("pytorch_wavelets")
+    torch.set_grad_enabled(
+        True
+    )  # Disabled somewhere in previous test files, necessary for this test to pass
+
+    # Select the data fidelity term
+    data_fidelity = L2()
+
+    # Set up the trainable denoising prior; here, the soft-threshold in a wavelet basis.
+    # If the prior is initialized with a list of length max_iter,
+    # then a distinct weight is trained for each PGD iteration.
+    # For fixed trained model prior across iterations, initialize with a single model.
+    max_iter = 30 if torch.cuda.is_available() else 20  # Number of unrolled iterations
+    level = 3
+    prior = [
+        PnP(denoiser=dinv.models.WaveletPrior(wv="db8", level=level, device=device))
+        for i in range(max_iter)
+    ]
+
+    # Unrolled optimization algorithm parameters
+    lamb = [
+        1.0
+    ] * max_iter  # initialization of the regularization parameter. A distinct lamb is trained for each iteration.
+    stepsize = [
+        1.0
+    ] * max_iter  # initialization of the stepsizes. A distinct stepsize is trained for each iteration.
+
+    sigma_denoiser_init = 0.01
+    sigma_denoiser = [sigma_denoiser_init * torch.ones(level, 3)] * max_iter
+    # sigma_denoiser = [torch.Tensor([sigma_denoiser_init])]*max_iter
+    params_algo = {  # wrap all the restoration parameters in a 'params_algo' dictionary
+        "stepsize": stepsize,
+        "g_param": sigma_denoiser,
+        "lambda": lamb,
+    }
+
+    trainable_params = [
+        "g_param",
+        "stepsize",
+    ]  # define which parameters from 'params_algo' are trainable
+
+    # Define the unfolded trainable model.
+    for and_acc in [False, True]:
+        # DRS, ADMM and CP algorithms are not real fixed-point algorithms on the primal variable
+
+        model = DEQ_builder(
+            unfolded_algo,
+            params_algo=params_algo,
+            trainable_params=trainable_params,
+            data_fidelity=data_fidelity,
+            max_iter=max_iter,
+            prior=prior,
+            anderson_acceleration=and_acc,
+            anderson_acceleration_backward=and_acc,
+        )
+
+        for idx, (name, param) in enumerate(model.named_parameters()):
+            assert param.requires_grad
+            assert (trainable_params[0] in name) or (trainable_params[1] in name)
+
+        # batch_size, n_channels, img_size_w, img_size_h = 5, imsize
+        batch_size = 5
+        n_channels, img_size_w, img_size_h = imsize
+        noise_level = 0.01
+
+        torch.manual_seed(0)
+        test_sample = torch.randn(batch_size, n_channels, img_size_w, img_size_h).to(
+            device
+        )
+        groundtruth_sample = torch.randn(
+            batch_size, n_channels, img_size_w, img_size_h
+        ).to(device)
+
+        physics = dinv.physics.BlurFFT(
+            img_size=(n_channels, img_size_w, img_size_h),
+            filter=dinv.physics.blur.gaussian_blur(),
+            device=device,
+            noise_model=dinv.physics.GaussianNoise(sigma=noise_level),
+        )
+
+        y = physics(test_sample).type(test_sample.dtype).to(device)
+
+        out = model(y, physics=physics)
+
+        assert out.shape == test_sample.shape
+
+        loss_fn = dinv.loss.SupLoss(metric=dinv.metric.mse())
+        loss = loss_fn(groundtruth_sample, out)
+        loss.backward()

deepinv/tests/test_utils.py

@@ -0,0 +1,68 @@
+import deepinv
+import torch
+import pytest
+
+
+@pytest.fixture
+def tensorlist():
+    x = torch.ones((1, 1, 2, 2))
+    y = torch.ones((1, 1, 2, 2))
+    x = deepinv.utils.TensorList([x, x])
+    y = deepinv.utils.TensorList([y, y])
+    return x, y
+
+
+def test_tensordict_sum(tensorlist):
+    x, y = tensorlist
+    z = torch.ones((1, 1, 2, 2)) * 2
+    z1 = deepinv.utils.TensorList([z, z])
+    z = x + y
+    assert (z1[0] == z[0]).all() and (z1[1] == z[1]).all()
+
+
+def test_tensordict_mul(tensorlist):
+    x, y = tensorlist
+    z = torch.ones((1, 1, 2, 2))
+    z1 = deepinv.utils.TensorList([z, z])
+    z = x * y
+    assert (z1[0] == z[0]).all() and (z1[1] == z[1]).all()
+
+
+def test_tensordict_div(tensorlist):
+    x, y = tensorlist
+    z = torch.ones((1, 1, 2, 2))
+    z1 = deepinv.utils.TensorList([z, z])
+    z = x / y
+    assert (z1[0] == z[0]).all() and (z1[1] == z[1]).all()
+
+
+def test_tensordict_sub(tensorlist):
+    x, y = tensorlist
+    z = torch.zeros((1, 1, 2, 2))
+    z1 = deepinv.utils.TensorList([z, z])
+    z = x - y
+    assert (z1[0] == z[0]).all() and (z1[1] == z[1]).all()
+
+
+def test_tensordict_neg(tensorlist):
+    x, y = tensorlist
+    z = -torch.ones((1, 1, 2, 2))
+    z1 = deepinv.utils.TensorList([z, z])
+    z = -x
+    assert (z1[0] == z[0]).all() and (z1[1] == z[1]).all()
+
+
+def test_tensordict_append(tensorlist):
+    x, y = tensorlist
+    z = torch.ones((1, 1, 2, 2))
+    z1 = deepinv.utils.TensorList([z, z, z, z])
+    z = x.append(y)
+    assert (z1[0] == z[0]).all() and (z1[-1] == z[-1]).all()
+
+
+def test_plot():
+    x = torch.ones((1, 1, 2, 2))
+    imgs = [x, x]
+    deepinv.utils.plot(imgs, titles=["a", "b"])
+    deepinv.utils.plot(x, titles="a")
+    deepinv.utils.plot(imgs)