careamics 0.1.0rc6__py3-none-any.whl → 0.1.0rc8__py3-none-any.whl

careamics/models/lvae/layers.py

@@ -0,0 +1,1998 @@
+"""
+Script containing the common basic blocks (nn.Module) reused by the LadderVAE architecture.
+
+Hierarchy in the model blocks:
+
+"""
+
+from copy import deepcopy
+from typing import Callable, Dict, Iterable, Literal, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torchvision.transforms.functional as F
+from torch.distributions import kl_divergence
+from torch.distributions.normal import Normal
+
+from .utils import (
+    StableLogVar,
+    StableMean,
+    crop_img_tensor,
+    kl_normal_mc,
+    pad_img_tensor,
+)
+
+
+class ResidualBlock(nn.Module):
+    """
+    Residual block with 2 convolutional layers.
+
+    Some architectural notes:
+        - The number of input, intermediate, and output channels is the same,
+        - Padding is always 'same',
+        - The 2 convolutional layers have the same groups,
+        - No stride allowed,
+        - Kernel sizes must be odd.
+
+    The output isgiven by: `out = gate(f(x)) + x`.
+    The presence of the gating mechanism is optional, and f(x) has different
+    structures depending on the `block_type` argument.
+    Specifically, `block_type` is a string specifying the block's structure, with:
+        a = activation
+        b = batch norm
+        c = conv layer
+        d = dropout.
+    For example, "bacdbacd" defines a block with 2x[batchnorm, activation, conv, dropout].
+    """
+
+    default_kernel_size = (3, 3)
+
+    def __init__(
+        self,
+        channels: int,
+        nonlin: Callable,
+        kernel: Union[int, Iterable[int]] = None,
+        groups: int = 1,
+        batchnorm: bool = True,
+        block_type: str = None,
+        dropout: float = None,
+        gated: bool = None,
+        skip_padding: bool = False,
+        conv2d_bias: bool = True,
+    ):
+        """
+        Constructor.
+
+        Parameters
+        ----------
+        channels: int
+            The number of input and output channels (they are the same).
+        nonlin: Callable
+            The non-linearity function used in the block (e.g., `nn.ReLU`).
+        kernel: Union[int, Iterable[int]], optional
+            The kernel size used in the convolutions of the block.
+            It can be either a single integer or a pair of integers defining the squared kernel.
+            Default is `None`.
+        groups: int, optional
+            The number of groups to consider in the convolutions. Default is 1.
+        batchnorm: bool, optional
+            Whether to use batchnorm layers. Default is `True`.
+        block_type: str, optional
+            A string specifying the block structure, check class docstring for more info.
+            Default is `None`.
+        dropout: float, optional
+            The dropout probability in dropout layers. If `None` dropout is not used.
+            Default is `None`.
+        gated: bool, optional
+            Whether to use gated layer. Default is `None`.
+        skip_padding: bool, optional
+            Whether to skip padding in convolutions. Default is `False`.
+        conv2d_bias: bool, optional
+            Whether to use bias term in convolutions. Default is `True`.
+        """
+        super().__init__()
+
+        # Set kernel size & padding
+        if kernel is None:
+            kernel = self.default_kernel_size
+        elif isinstance(kernel, int):
+            kernel = (kernel, kernel)
+        elif len(kernel) != 2:
+            raise ValueError("kernel has to be None, int, or an iterable of length 2")
+        assert all([k % 2 == 1 for k in kernel]), "kernel sizes have to be odd"
+        kernel = list(kernel)
+        self.skip_padding = skip_padding
+        pad = [0] * len(kernel) if self.skip_padding else [k // 2 for k in kernel]
+        # print(kernel, pad)
+
+        modules = []
+        if block_type == "cabdcabd":
+            for i in range(2):
+                conv = nn.Conv2d(
+                    channels,
+                    channels,
+                    kernel[i],
+                    padding=pad[i],
+                    groups=groups,
+                    bias=conv2d_bias,
+                )
+                modules.append(conv)
+                modules.append(nonlin())
+                if batchnorm:
+                    modules.append(nn.BatchNorm2d(channels))
+                if dropout is not None:
+                    modules.append(nn.Dropout2d(dropout))
+        elif block_type == "bacdbac":
+            for i in range(2):
+                if batchnorm:
+                    modules.append(nn.BatchNorm2d(channels))
+                modules.append(nonlin())
+                conv = nn.Conv2d(
+                    channels,
+                    channels,
+                    kernel[i],
+                    padding=pad[i],
+                    groups=groups,
+                    bias=conv2d_bias,
+                )
+                modules.append(conv)
+                if dropout is not None and i == 0:
+                    modules.append(nn.Dropout2d(dropout))
+        elif block_type == "bacdbacd":
+            for i in range(2):
+                if batchnorm:
+                    modules.append(nn.BatchNorm2d(channels))
+                modules.append(nonlin())
+                conv = nn.Conv2d(
+                    channels,
+                    channels,
+                    kernel[i],
+                    padding=pad[i],
+                    groups=groups,
+                    bias=conv2d_bias,
+                )
+                modules.append(conv)
+                modules.append(nn.Dropout2d(dropout))
+
+        else:
+            raise ValueError(f"unrecognized block type '{block_type}'")
+
+        self.gated = gated
+        if gated:
+            modules.append(GateLayer2d(channels, 1, nonlin))
+
+        self.block = nn.Sequential(*modules)
+
+    def forward(self, x):
+
+        out = self.block(x)
+        if out.shape != x.shape:
+            return out + F.center_crop(x, out.shape[-2:])
+        else:
+            return out + x
+
+
+class ResidualGatedBlock(ResidualBlock):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs, gated=True)
+
+
+class GateLayer2d(nn.Module):
+    """
+    Double the number of channels through a convolutional layer, then use
+    half the channels as gate for the other half.
+    """
+
+    def __init__(self, channels, kernel_size, nonlin=nn.LeakyReLU):
+        super().__init__()
+        assert kernel_size % 2 == 1
+        pad = kernel_size // 2
+        self.conv = nn.Conv2d(channels, 2 * channels, kernel_size, padding=pad)
+        self.nonlin = nonlin()
+
+    def forward(self, x):
+        x = self.conv(x)
+        x, gate = torch.chunk(x, 2, dim=1)
+        x = self.nonlin(x)  # TODO remove this?
+        gate = torch.sigmoid(gate)
+        return x * gate
+
+
+class ResBlockWithResampling(nn.Module):
+    """
+    Residual block that takes care of resampling (i.e. downsampling or upsampling) steps (by a factor 2).
+    It is structured as follows:
+        1. `pre_conv`: a downsampling or upsampling strided convolutional layer in case of resampling, or
+            a 1x1 convolutional layer that maps the number of channels of the input to `inner_channels`.
+        2. `ResidualBlock`
+        3. `post_conv`: a 1x1 convolutional layer that maps the number of channels to `c_out`.
+
+    Some implementation notes:
+    - Resampling is performed through a strided convolution layer at the beginning of the block.
+    - The strided convolution block has fixed kernel size of 3x3 and 1 layer of zero-padding.
+    - The number of channels is adjusted at the beginning and end of the block through 1x1 convolutional layers.
+    - The number of internal channels is by default the same as the number of output channels, but
+      min_inner_channels can override the behaviour.
+    """
+
+    def __init__(
+        self,
+        mode: Literal["top-down", "bottom-up"],
+        c_in: int,
+        c_out: int,
+        min_inner_channels: int = None,
+        nonlin: Callable = nn.LeakyReLU,
+        resample: bool = False,
+        res_block_kernel: Union[int, Iterable[int]] = None,
+        groups: int = 1,
+        batchnorm: bool = True,
+        res_block_type: str = None,
+        dropout: float = None,
+        gated: bool = None,
+        skip_padding: bool = False,
+        conv2d_bias: bool = True,
+        # lowres_input: bool = False,
+    ):
+        """
+        Constructor.
+
+        Parameters
+        ----------
+        mode: Literal["top-down", "bottom-up"]
+            The type of resampling performed in the initial strided convolution of the block.
+            If "bottom-up" downsampling of a factor 2 is done.
+            If "top-down" upsampling of a factor 2 is done.
+        c_in: int
+            The number of input channels.
+        c_out: int
+            The number of output channels.
+        min_inner_channels: int, optional
+            The number of channels used in the inner layer of this module.
+            Default is `None`, meaning that the number of inner channels is set to `c_out`.
+        nonlin: Callable, optional
+            The non-linearity function used in the block. Default is `nn.LeakyReLU`.
+        resample: bool, optional
+            Whether to perform resampling in the first convolutional layer.
+            If `False`, the first convolutional layer just maps the input to a tensor with
+            `inner_channels` channels through 1x1 convolution. Deafult is `False`.
+        res_block_kernel: Union[int, Iterable[int]], optional
+            The kernel size used in the convolutions of the residual block.
+            It can be either a single integer or a pair of integers defining the squared kernel.
+            Default is `None`.
+        groups: int, optional
+            The number of groups to consider in the convolutions. Default is 1.
+        batchnorm: bool, optional
+            Whether to use batchnorm layers. Default is `True`.
+        res_block_type: str, optional
+            A string specifying the structure of residual block.
+            Check `ResidualBlock` doscstring for more information.
+            Default is `None`.
+        dropout: float, optional
+            The dropout probability in dropout layers. If `None` dropout is not used.
+            Default is `None`.
+        gated: bool, optional
+            Whether to use gated layer. Default is `None`.
+        skip_padding: bool, optional
+            Whether to skip padding in convolutions. Default is `False`.
+        conv2d_bias: bool, optional
+            Whether to use bias term in convolutions. Default is `True`.
+        """
+        super().__init__()
+        assert mode in ["top-down", "bottom-up"]
+
+        if min_inner_channels is None:
+            min_inner_channels = 0
+        # inner_channels is the number of channels used in the inner layers
+        # of ResBlockWithResampling
+        inner_channels = max(c_out, min_inner_channels)
+
+        # Define first conv layer to change num channels and/or up/downsample
+        if resample:
+            if mode == "bottom-up":  # downsample
+                self.pre_conv = nn.Conv2d(
+                    in_channels=c_in,
+                    out_channels=inner_channels,
+                    kernel_size=3,
+                    padding=1,
+                    stride=2,
+                    groups=groups,
+                    bias=conv2d_bias,
+                )
+            elif mode == "top-down":  # upsample
+                self.pre_conv = nn.ConvTranspose2d(
+                    in_channels=c_in,
+                    kernel_size=3,
+                    out_channels=inner_channels,
+                    padding=1,
+                    stride=2,
+                    groups=groups,
+                    output_padding=1,
+                    bias=conv2d_bias,
+                )
+        elif c_in != inner_channels:
+            self.pre_conv = nn.Conv2d(
+                c_in, inner_channels, 1, groups=groups, bias=conv2d_bias
+            )
+        else:
+            self.pre_conv = None
+
+        # Residual block
+        self.res = ResidualBlock(
+            channels=inner_channels,
+            nonlin=nonlin,
+            kernel=res_block_kernel,
+            groups=groups,
+            batchnorm=batchnorm,
+            dropout=dropout,
+            gated=gated,
+            block_type=res_block_type,
+            skip_padding=skip_padding,
+            conv2d_bias=conv2d_bias,
+        )
+
+        # Define last conv layer to get correct num output channels
+        if inner_channels != c_out:
+            self.post_conv = nn.Conv2d(
+                inner_channels, c_out, 1, groups=groups, bias=conv2d_bias
+            )
+        else:
+            self.post_conv = None
+
+    def forward(self, x):
+        if self.pre_conv is not None:
+            x = self.pre_conv(x)
+
+        x = self.res(x)
+
+        if self.post_conv is not None:
+            x = self.post_conv(x)
+        return x
+
+
+class TopDownDeterministicResBlock(ResBlockWithResampling):
+
+    def __init__(self, *args, upsample: bool = False, **kwargs):
+        kwargs["resample"] = upsample
+        super().__init__("top-down", *args, **kwargs)
+
+
+class BottomUpDeterministicResBlock(ResBlockWithResampling):
+
+    def __init__(self, *args, downsample: bool = False, **kwargs):
+        kwargs["resample"] = downsample
+        super().__init__("bottom-up", *args, **kwargs)
+
+
+class BottomUpLayer(nn.Module):
+    """
+    Bottom-up deterministic layer.
+    It consists of one or a stack of `BottomUpDeterministicResBlock`'s.
+    The outputs are the so-called `bu_values` that are later used in the Decoder to update the
+    generative distributions.
+
+    NOTE: When Lateral Contextualization is Enabled (i.e., `enable_multiscale=True`),
+    the low-res lateral input is first fed through a BottomUpDeterministicBlock (BUDB)
+    (without downsampling), and then merged to the latent tensor produced by the primary flow
+    of the `BottomUpLayer` through the `MergeLowRes` layer. It is meaningful to remark that
+    the BUDB that takes care of encoding the low-res input can be either shared with the
+    primary flow (and in that case it is the "same_size" BUDB (or stack of BUDBs) -> see `self.net`),
+    or can be a deep-copy of the primary flow's BUDB.
+    This behaviour is controlled by `lowres_separate_branch` parameter.
+    """
+
+    def __init__(
+        self,
+        n_res_blocks: int,
+        n_filters: int,
+        downsampling_steps: int = 0,
+        nonlin: Callable = None,
+        batchnorm: bool = True,
+        dropout: float = None,
+        res_block_type: str = None,
+        res_block_kernel: int = None,
+        res_block_skip_padding: bool = False,
+        gated: bool = None,
+        enable_multiscale: bool = False,
+        multiscale_lowres_size_factor: int = None,
+        lowres_separate_branch: bool = False,
+        multiscale_retain_spatial_dims: bool = False,
+        decoder_retain_spatial_dims: bool = False,
+        output_expected_shape: Iterable[int] = None,
+    ):
+        """
+        Constructor.
+
+        Parameters
+        ----------
+        n_res_blocks: int
+            Number of `BottomUpDeterministicResBlock` modules stacked in this layer.
+        n_filters: int
+            Number of channels present through out the layers of this block.
+        downsampling_steps: int, optional
+            Number of downsampling steps that has to be done in this layer (typically 1).
+            Default is 0.
+        nonlin: Callable, optional
+            The non-linearity function used in the block. Default is `None`.
+        batchnorm: bool, optional
+            Whether to use batchnorm layers. Default is `True`.
+        dropout: float, optional
+            The dropout probability in dropout layers. If `None` dropout is not used.
+            Default is `None`.
+        res_block_type: str, optional
+            A string specifying the structure of residual block.
+            Check `ResidualBlock` doscstring for more information.
+            Default is `None`.
+        res_block_kernel: Union[int, Iterable[int]], optional
+            The kernel size used in the convolutions of the residual block.
+            It can be either a single integer or a pair of integers defining the squared kernel.
+            Default is `None`.
+        res_block_skip_padding: bool, optional
+            Whether to skip padding in convolutions in the Residual block. Default is `False`.
+        gated: bool, optional
+            Whether to use gated layer. Default is `None`.
+        enable_multiscale: bool, optional
+            Whether to enable multiscale (Lateral Contextualization) or not. Default is `False`.
+        multiscale_lowres_size_factor: int, optional
+            A factor the expresses the relative size of the primary flow tensor with respect to the
+            lower-resolution lateral input tensor. Default in `None`.
+        lowres_separate_branch: bool, optional
+            Whether the residual block(s) encoding the low-res input should be shared (`False`) or
+            not (`True`) with the primary flow "same-size" residual block(s). Default is `False`.
+        multiscale_retain_spatial_dims: bool, optional
+            Whether to pad the latent tensor resulting from the bottom-up layer's primary flow
+            to match the size of the low-res input. Default is `False`.
+        decoder_retain_spatial_dims: bool, optional
+            Default is `False`.
+        output_expected_shape: Iterable[int], optional
+            The expected shape of the layer output (only used if `enable_multiscale == True`).
+            Default is `None`.
+        """
+        super().__init__()
+
+        # Define attributes for Lateral Contextualization
+        self.enable_multiscale = enable_multiscale
+        self.lowres_separate_branch = lowres_separate_branch
+        self.multiscale_retain_spatial_dims = multiscale_retain_spatial_dims
+        self.multiscale_lowres_size_factor = multiscale_lowres_size_factor
+        self.decoder_retain_spatial_dims = decoder_retain_spatial_dims
+        self.output_expected_shape = output_expected_shape
+        assert self.output_expected_shape is None or self.enable_multiscale is True
+
+        bu_blocks_downsized = []
+        bu_blocks_samesize = []
+        for _ in range(n_res_blocks):
+            do_resample = False
+            if downsampling_steps > 0:
+                do_resample = True
+                downsampling_steps -= 1
+            block = BottomUpDeterministicResBlock(
+                c_in=n_filters,
+                c_out=n_filters,
+                nonlin=nonlin,
+                downsample=do_resample,
+                batchnorm=batchnorm,
+                dropout=dropout,
+                res_block_type=res_block_type,
+                res_block_kernel=res_block_kernel,
+                skip_padding=res_block_skip_padding,
+                gated=gated,
+            )
+            if do_resample:
+                bu_blocks_downsized.append(block)
+            else:
+                bu_blocks_samesize.append(block)
+
+        self.net_downsized = nn.Sequential(*bu_blocks_downsized)
+        self.net = nn.Sequential(*bu_blocks_samesize)
+
+        # Using the same net for the low resolution (and larger sized image)
+        self.lowres_net = self.lowres_merge = None
+        if self.enable_multiscale:
+            self._init_multiscale(
+                n_filters=n_filters,
+                nonlin=nonlin,
+                batchnorm=batchnorm,
+                dropout=dropout,
+                res_block_type=res_block_type,
+            )
+
+        # msg = f'[{self.__class__.__name__}] McEnabled:{int(enable_multiscale)} '
+        # if enable_multiscale:
+        #     msg += f'McParallelBeam:{int(multiscale_retain_spatial_dims)} McFactor{multiscale_lowres_size_factor}'
+        # print(msg)
+
+    def _init_multiscale(
+        self,
+        nonlin: Callable = None,
+        n_filters: int = None,
+        batchnorm: bool = None,
+        dropout: float = None,
+        res_block_type: str = None,
+    ) -> None:
+        """
+        This method defines the modules responsible of merging compressed lateral inputs to the outputs
+        of the primary flow at different hierarchical levels in the multiresolution approach (LC).
+
+        Specifically, the method initializes `lowres_net`, which is a stack of `BottomUpDeterministicBlock`'s
+        (w/out downsampling) that takes care of additionally processing the low-res input, and `lowres_merge`,
+        which is the module responsible of merging the compressed lateral input to the main flow.
+
+        NOTE: The merge modality is set by default to "residual", meaning that the merge layer
+        performs concatenation on dim=1, followed by 1x1 convolution and a Residual Gated block.
+
+        Parameters
+        ----------
+        nonlin: Callable, optional
+            The non-linearity function used in the block. Default is `None`.
+        n_filters: int
+            Number of channels present through out the layers of this block.
+        batchnorm: bool, optional
+            Whether to use batchnorm layers. Default is `True`.
+        dropout: float, optional
+            The dropout probability in dropout layers. If `None` dropout is not used.
+            Default is `None`.
+        res_block_type: str, optional
+            A string specifying the structure of residual block.
+            Check `ResidualBlock` doscstring for more information.
+            Default is `None`.
+        """
+        self.lowres_net = self.net
+        if self.lowres_separate_branch:
+            self.lowres_net = deepcopy(self.net)
+
+        self.lowres_merge = MergeLowRes(
+            channels=n_filters,
+            merge_type="residual",
+            nonlin=nonlin,
+            batchnorm=batchnorm,
+            dropout=dropout,
+            res_block_type=res_block_type,
+            multiscale_retain_spatial_dims=self.multiscale_retain_spatial_dims,
+            multiscale_lowres_size_factor=self.multiscale_lowres_size_factor,
+        )
+
+    def forward(
+        self, x: torch.Tensor, lowres_x: torch.Tensor = None
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Parameters
+        ----------
+        x: torch.Tensor
+            The input of the `BottomUpLayer`, i.e., the input image or the output of the
+            previous layer.
+        lowres_x: torch.Tensor, optional
+            The low-res input used for Lateral Contextualization (LC). Default is `None`.
+        """
+        # The input is fed through the residual downsampling block(s)
+        primary_flow = self.net_downsized(x)
+        # The downsampling output is fed through additional residual block(s)
+        primary_flow = self.net(primary_flow)
+
+        # If LC is not used, simply return output of primary-flow
+        if self.enable_multiscale is False:
+            assert lowres_x is None
+            return primary_flow, primary_flow
+
+        if lowres_x is not None:
+            # First encode the low-res lateral input
+            lowres_flow = self.lowres_net(lowres_x)
+            # Then pass the result through the MergeLowRes layer
+            merged = self.lowres_merge(primary_flow, lowres_flow)
+        else:
+            merged = primary_flow
+
+        if (
+            self.multiscale_retain_spatial_dims is False
+            or self.decoder_retain_spatial_dims is True
+        ):
+            return merged, merged
+
+        if self.output_expected_shape is not None:
+            expected_shape = self.output_expected_shape
+        else:
+            fac = self.multiscale_lowres_size_factor
+            expected_shape = (merged.shape[-2] // fac, merged.shape[-1] // fac)
+            assert merged.shape[-2:] != expected_shape
+
+        # Crop the resulting tensor so that it matches with the Decoder
+        value_to_use_in_topdown = crop_img_tensor(merged, expected_shape)
+        return merged, value_to_use_in_topdown
+
+
+class MergeLayer(nn.Module):
+    """
+    This layer merges two or more 4D input tensors by concatenating along dim=1 and passes the result through:
+    a) a convolutional 1x1 layer (`merge_type == "linear"`), or
+    b) a convolutional 1x1 layer and then a gated residual block (`merge_type == "residual"`), or
+    c) a convolutional 1x1 layer and then an ungated residual block (`merge_type == "residual_ungated"`).
+    """
+
+    def __init__(
+        self,
+        merge_type: Literal["linear", "residual", "residual_ungated"],
+        channels: Union[int, Iterable[int]],
+        nonlin: Callable = nn.LeakyReLU,
+        batchnorm: bool = True,
+        dropout: float = None,
+        res_block_type: str = None,
+        res_block_kernel: int = None,
+        res_block_skip_padding: bool = False,
+        conv2d_bias: bool = True,
+    ):
+        """
+        Constructor.
+
+        Parameters
+        ----------
+        merge_type: Literal["linear", "residual", "residual_ungated"]
+            The type of merge done in the layer. It can be chosen between "linear", "residual", and "residual_ungated".
+            Check the class docstring for more information about the behaviour of different merge modalities.
+        channels: Union[int, Iterable[int]]
+            The number of channels used in the convolutional blocks of this layer.
+            If it is an `int`:
+                - 1st 1x1 Conv2d: in_channels=2*channels, out_channels=channels
+                - (Optional) ResBlock: in_channels=channels, out_channels=channels
+            If it is an Iterable (must have `len(channels)==3`):
+                - 1st 1x1 Conv2d: in_channels=sum(channels[:-1]), out_channels=channels[-1]
+                - (Optional) ResBlock: in_channels=channels[-1], out_channels=channels[-1]
+        nonlin: Callable, optional
+            The non-linearity function used in the block. Default is `nn.LeakyReLU`.
+        batchnorm: bool, optional
+            Whether to use batchnorm layers. Default is `True`.
+        dropout: float, optional
+            The dropout probability in dropout layers. If `None` dropout is not used.
+            Default is `None`.
+        res_block_type: str, optional
+            A string specifying the structure of residual block.
+            Check `ResidualBlock` doscstring for more information.
+            Default is `None`.
+        res_block_kernel: Union[int, Iterable[int]], optional
+            The kernel size used in the convolutions of the residual block.
+            It can be either a single integer or a pair of integers defining the squared kernel.
+            Default is `None`.
+        res_block_skip_padding: bool, optional
+            Whether to skip padding in convolutions in the Residual block. Default is `False`.
+        conv2d_bias: bool, optional
+            Whether to use bias term in convolutions. Default is `True`.
+        """
+        super().__init__()
+        try:
+            iter(channels)
+        except TypeError:  # it is not iterable
+            channels = [channels] * 3
+        else:  # it is iterable
+            if len(channels) == 1:
+                channels = [channels[0]] * 3
+
+        # assert len(channels) == 3
+
+        if merge_type == "linear":
+            self.layer = nn.Conv2d(
+                sum(channels[:-1]), channels[-1], 1, bias=conv2d_bias
+            )
+        elif merge_type == "residual":
+            self.layer = nn.Sequential(
+                nn.Conv2d(
+                    sum(channels[:-1]), channels[-1], 1, padding=0, bias=conv2d_bias
+                ),
+                ResidualGatedBlock(
+                    channels[-1],
+                    nonlin,
+                    batchnorm=batchnorm,
+                    dropout=dropout,
+                    block_type=res_block_type,
+                    kernel=res_block_kernel,
+                    conv2d_bias=conv2d_bias,
+                    skip_padding=res_block_skip_padding,
+                ),
+            )
+        elif merge_type == "residual_ungated":
+            self.layer = nn.Sequential(
+                nn.Conv2d(
+                    sum(channels[:-1]), channels[-1], 1, padding=0, bias=conv2d_bias
+                ),
+                ResidualBlock(
+                    channels[-1],
+                    nonlin,
+                    batchnorm=batchnorm,
+                    dropout=dropout,
+                    block_type=res_block_type,
+                    kernel=res_block_kernel,
+                    conv2d_bias=conv2d_bias,
+                    skip_padding=res_block_skip_padding,
+                ),
+            )
+
+    def forward(self, *args) -> torch.Tensor:
+
+        # Concatenate the input tensors along dim=1
+        x = torch.cat(args, dim=1)
+
+        # Pass the concatenated tensor through the conv layer
+        x = self.layer(x)
+
+        return x
+
+
+class MergeLowRes(MergeLayer):
+    """
+    Child class of `MergeLayer`, specifically designed to merge the low-resolution patches
+    that are used in Lateral Contextualization approach.
+    """
+
+    def __init__(self, *args, **kwargs):
+        self.retain_spatial_dims = kwargs.pop("multiscale_retain_spatial_dims")
+        self.multiscale_lowres_size_factor = kwargs.pop("multiscale_lowres_size_factor")
+        super().__init__(*args, **kwargs)
+
+    def forward(self, latent: torch.Tensor, lowres: torch.Tensor) -> torch.Tensor:
+        """
+        Parameters
+        ----------
+        latent: torch.Tensor
+            The output latent tensor from previous layer in the LVAE hierarchy.
+        lowres: torch.Tensor
+            The low-res patch image to be merged to increase the context.
+        """
+        if self.retain_spatial_dims:
+            # Pad latent tensor to match lowres tensor's shape
+            latent = pad_img_tensor(latent, lowres.shape[2:])
+        else:
+            # Crop lowres tensor to match latent tensor's shape
+            lh, lw = lowres.shape[-2:]
+            h = lh // self.multiscale_lowres_size_factor
+            w = lw // self.multiscale_lowres_size_factor
+            h_pad = (lh - h) // 2
+            w_pad = (lw - w) // 2
+            lowres = lowres[:, :, h_pad:-h_pad, w_pad:-w_pad]
+
+        return super().forward(latent, lowres)
+
+
+class SkipConnectionMerger(MergeLayer):
+    """
+    A specialized `MergeLayer` module, designed to handle skip connections in the model.
+    """
+
+    def __init__(
+        self,
+        nonlin: Callable,
+        channels: Union[int, Iterable[int]],
+        batchnorm: bool,
+        dropout: float,
+        res_block_type: str,
+        merge_type: Literal["linear", "residual", "residual_ungated"] = "residual",
+        conv2d_bias: bool = True,
+        res_block_kernel: int = None,
+        res_block_skip_padding: bool = False,
+    ):
+        """
+        Constructor.
+
+        nonlin: Callable, optional
+            The non-linearity function used in the block. Default is `nn.LeakyReLU`.
+        channels: Union[int, Iterable[int]]
+            The number of channels used in the convolutional blocks of this layer.
+            If it is an `int`:
+                - 1st 1x1 Conv2d: in_channels=2*channels, out_channels=channels
+                - (Optional) ResBlock: in_channels=channels, out_channels=channels
+            If it is an Iterable (must have `len(channels)==3`):
+                - 1st 1x1 Conv2d: in_channels=sum(channels[:-1]), out_channels=channels[-1]
+                - (Optional) ResBlock: in_channels=channels[-1], out_channels=channels[-1]
+        batchnorm: bool, optional
+            Whether to use batchnorm layers. Default is `True`.
+        dropout: float, optional
+            The dropout probability in dropout layers. If `None` dropout is not used.
+            Default is `None`.
+        res_block_type: str, optional
+            A string specifying the structure of residual block.
+            Check `ResidualBlock` doscstring for more information.
+            Default is `None`.
+        merge_type: Literal["linear", "residual", "residual_ungated"]
+            The type of merge done in the layer. It can be chosen between "linear", "residual", and "residual_ungated".
+            Check the class docstring for more information about the behaviour of different merge modalities.
+        conv2d_bias: bool, optional
+            Whether to use bias term in convolutions. Default is `True`.
+        res_block_kernel: Union[int, Iterable[int]], optional
+            The kernel size used in the convolutions of the residual block.
+            It can be either a single integer or a pair of integers defining the squared kernel.
+            Default is `None`.
+        res_block_skip_padding: bool, optional
+            Whether to skip padding in convolutions in the Residual block. Default is `False`.
+        """
+        super().__init__(
+            channels=channels,
+            nonlin=nonlin,
+            merge_type=merge_type,
+            batchnorm=batchnorm,
+            dropout=dropout,
+            res_block_type=res_block_type,
+            res_block_kernel=res_block_kernel,
+            conv2d_bias=conv2d_bias,
+            res_block_skip_padding=res_block_skip_padding,
+        )
+
+
+class TopDownLayer(nn.Module):
+    """
+    Top-down inference layer.
+    It includes:
+        - Stochastic sampling,
+        - Computation of KL divergence,
+        - A small deterministic ResNet that performs upsampling.
+
+    NOTE 1:
+            The algorithm for generative inference approximately works as follows:
+                - p_params = output of top-down layer above
+                - bu = inferred bottom-up value at this layer
+                - q_params = merge(bu, p_params)
+                - z = stochastic_layer(q_params)
+                - (optional) get and merge skip connection from prev top-down layer
+                - top-down deterministic ResNet
+
+    NOTE 2:
+        The Top-Down layer can work in two modes: inference and prediction/generative.
+        Depending on the particular mode, it follows distinct behaviours:
+        - In inference mode, parameters of q(z_i|z_i+1) are obtained from the inference path,
+        by merging outcomes of bottom-up and top-down passes. The exception is the top layer,
+        in which the parameters of q(z_L|x) are set as the output of the topmost bottom-up layer.
+        - On the contrary in prediciton/generative mode, parameters of q(z_i|z_i+1) can be obtained
+        once again by merging bottom-up and top-down outputs (CONDITIONAL GENERATION), or it is
+        possible to directly sample from the prior p(z_i|z_i+1) (UNCONDITIONAL GENERATION).
+
+    NOTE 3:
+        When doing unconditional generation, bu_value is not available. Hence the
+        merge layer is not used, and z is sampled directly from p_params.
+
+    NOTE 4:
+        If this is the top layer, at inference time, the uppermost bottom-up value
+        is used directly as q_params, and p_params are defined in this layer
+        (while they are usually taken from the previous layer), and can be learned.
+    """
+
+    def __init__(
+        self,
+        z_dim: int,
+        n_res_blocks: int,
+        n_filters: int,
+        is_top_layer: bool = False,
+        downsampling_steps: int = None,
+        nonlin: Callable = None,
+        merge_type: Literal["linear", "residual", "residual_ungated"] = None,
+        batchnorm: bool = True,
+        dropout: float = None,
+        stochastic_skip: bool = False,
+        res_block_type: str = None,
+        res_block_kernel: int = None,
+        res_block_skip_padding: bool = None,
+        groups: int = 1,
+        gated: bool = None,
+        learn_top_prior: bool = False,
+        top_prior_param_shape: Iterable[int] = None,
+        analytical_kl: bool = False,
+        bottomup_no_padding_mode: bool = False,
+        topdown_no_padding_mode: bool = False,
+        retain_spatial_dims: bool = False,
+        restricted_kl: bool = False,
+        vanilla_latent_hw: Iterable[int] = None,
+        non_stochastic_version: bool = False,
+        input_image_shape: Union[None, Tuple[int, int]] = None,
+        normalize_latent_factor: float = 1.0,
+        conv2d_bias: bool = True,
+        stochastic_use_naive_exponential: bool = False,
+    ):
+        """
+        Constructor.
+
+        Parameters
+        ----------
+        z_dim: int
+            The size of the latent space.
+        n_res_blocks: int
+            The number of TopDownDeterministicResBlock blocks
+        n_filters: int
+            The number of channels present through out the layers of this block.
+        is_top_layer: bool, optional
+            Whether the current layer is at the top of the Decoder hierarchy. Default is `False`.
+        downsampling_steps: int, optional
+            The number of downsampling steps that has to be done in this layer (typically 1).
+            Default is `False`.
+        nonlin: Callable, optional
+            The non-linearity function used in the block (e.g., `nn.ReLU`). Deafault is `None`.
+        merge_type: Literal["linear", "residual", "residual_ungated"], optional
+            The type of merge done in the layer. It can be chosen between "linear", "residual",
+            and "residual_ungated". Check the `MergeLayer` class docstring for more information
+            about the behaviour of different merging modalities. Default is `None`.
+        batchnorm: bool, optional
+            Whether to use batchnorm layers. Default is `True`.
+        dropout: float, optional
+            The dropout probability in dropout layers. If `None` dropout is not used.
+            Default is `None`.
+        stochastic_skip: bool, optional
+            Whether to use skip connections between previous top-down layer's output and this layer's stochastic output.
+            Stochastic skip connection allows the previous layer's output has a way to directly reach this hierarchical
+            level, hence facilitating the gradient flow during backpropagation. Default is `False`.
+        res_block_type: str, optional
+            A string specifying the structure of residual block.
+            Check `ResidualBlock` documentation for more information.
+            Default is `None`.
+        res_block_kernel: Union[int, Iterable[int]], optional
+            The kernel size used in the convolutions of the residual block.
+            It can be either a single integer or a pair of integers defining the squared kernel.
+            Default is `None`.
+        res_block_skip_padding: bool, optional
+            Whether to skip padding in convolutions in the Residual block. Default is `None`.
+        groups: int, optional
+            The number of groups to consider in the convolutions. Default is 1.
+        gated: bool, optional
+            Whether to use gated layer in `ResidualBlock`. Default is `None`.
+        learn_top_prior:
+            Whether to set the top prior as learnable.
+            If this is set to `False`, in the top-most layer the prior will be N(0,1).
+            Otherwise, we will still have a normal distribution whose parameters will be learnt.
+            Deafult is `False`.
+        top_prior_param_shape: Iterable[int], optional
+            The size of the tensor which expresses the mean and the variance
+            of the prior for the top most layer. Default is `None`.
+        analytical_kl: bool, optional
+            If True, KL divergence is calculated according to the analytical formula.
+            Otherwise, an MC approximation using sampled latents is calculated.
+            Default is `False`.
+        bottomup_no_padding_mode: bool, optional
+            Whether padding is used in the different layers of the bottom-up pass.
+            It is meaningful to know this in advance in order to assess whether before
+            merging `bu_values` and `p_params` tensors any alignment is needed.
+            Default is `False`.
+        topdown_no_padding_mode: bool, optional
+            Whether padding is used in the different layers of the top-down pass.
+            It is meaningful to know this in advance in order to assess whether before
+            merging `bu_values` and `p_params` tensors any alignment is needed.
+            The same information is also needed in handling the skip connections between
+            top-down layers. Default is `False`.
+        retain_spatial_dims: bool, optional
+            If `True`, the size of Encoder's latent space is kept to `input_image_shape` within the topdown layer.
+            This implies that the oput spatial size equals the input spatial size.
+            To achieve this, we centercrop the intermediate representation.
+            Default is `False`.
+        restricted_kl: bool, optional
+            Whether to compute the restricted version of KL Divergence.
+            See `NormalStochasticBlock2d` module for more information about its computation.
+            Default is `False`.
+        vanilla_latent_hw: Iterable[int], optional
+            The shape of the latent tensor used for prediction (i.e., it influences the computation of restricted KL).
+            Default is `None`.
+        non_stochastic_version: bool, optional
+            Whether to replace the stochastic layer that samples a latent variable from the latent distribiution with
+            a non-stochastic layer that simply drwas a sample as the mode of the latent distribution.
+            Default is `False`.
+        input_image_shape: Tuple[int, int], optionalut
+            The shape of the input image tensor.
+            When `retain_spatial_dims` is set to `True`, this is used to ensure that the shape of this layer
+            output has the same shape as the input. Default is `None`.
+        normalize_latent_factor: float, optional
+            A factor used to normalize the latent tensors `q_params`.
+            Specifically, normalization is done by dividing the latent tensor by this factor.
+            Default is 1.0.
+        conv2d_bias: bool, optional
+            Whether to use bias term is the convolutional blocks of this layer.
+            Default is `True`.
+        stochastic_use_naive_exponential: bool, optional
+            If `False`, in the NormalStochasticBlock2d exponentials are computed according
+            to the alternative definition provided by `StableExponential` class.
+            This should improve numerical stability in the training process.
+            Default is `False`.
+        """
+        super().__init__()
+
+        self.is_top_layer = is_top_layer
+        self.z_dim = z_dim
+        self.stochastic_skip = stochastic_skip
+        self.learn_top_prior = learn_top_prior
+        self.analytical_kl = analytical_kl
+        self.bottomup_no_padding_mode = bottomup_no_padding_mode
+        self.topdown_no_padding_mode = topdown_no_padding_mode
+        self.retain_spatial_dims = retain_spatial_dims
+        self.latent_shape = input_image_shape if self.retain_spatial_dims else None
+        self.non_stochastic_version = non_stochastic_version
+        self.normalize_latent_factor = normalize_latent_factor
+        self._vanilla_latent_hw = vanilla_latent_hw
+
+        # Define top layer prior parameters, possibly learnable
+        if is_top_layer:
+            self.top_prior_params = nn.Parameter(
+                torch.zeros(top_prior_param_shape), requires_grad=learn_top_prior
+            )
+
+        # Downsampling steps left to do in this layer
+        dws_left = downsampling_steps
+
+        # Define deterministic top-down block, which is a sequence of deterministic
+        # residual blocks with (optional) downsampling.
+        block_list = []
+        for _ in range(n_res_blocks):
+            do_resample = False
+            if dws_left > 0:
+                do_resample = True
+                dws_left -= 1
+            block_list.append(
+                TopDownDeterministicResBlock(
+                    c_in=n_filters,
+                    c_out=n_filters,
+                    nonlin=nonlin,
+                    upsample=do_resample,
+                    batchnorm=batchnorm,
+                    dropout=dropout,
+                    res_block_type=res_block_type,
+                    res_block_kernel=res_block_kernel,
+                    skip_padding=res_block_skip_padding,
+                    gated=gated,
+                    conv2d_bias=conv2d_bias,
+                    groups=groups,
+                )
+            )
+        self.deterministic_block = nn.Sequential(*block_list)
+
+        # Define stochastic block with 2D convolutions
+        if self.non_stochastic_version:
+            self.stochastic = NonStochasticBlock2d(
+                c_in=n_filters,
+                c_vars=z_dim,
+                c_out=n_filters,
+                transform_p_params=(not is_top_layer),
+                groups=groups,
+                conv2d_bias=conv2d_bias,
+            )
+        else:
+            self.stochastic = NormalStochasticBlock2d(
+                c_in=n_filters,
+                c_vars=z_dim,
+                c_out=n_filters,
+                transform_p_params=(not is_top_layer),
+                vanilla_latent_hw=vanilla_latent_hw,
+                restricted_kl=restricted_kl,
+                use_naive_exponential=stochastic_use_naive_exponential,
+            )
+
+        if not is_top_layer:
+            # Merge layer: it combines bottom-up inference and top-down
+            # generative outcomes to give posterior parameters
+            self.merge = MergeLayer(
+                channels=n_filters,
+                merge_type=merge_type,
+                nonlin=nonlin,
+                batchnorm=batchnorm,
+                dropout=dropout,
+                res_block_type=res_block_type,
+                res_block_kernel=res_block_kernel,
+                conv2d_bias=conv2d_bias,
+            )
+
+            # Skip connection that goes around the stochastic top-down layer
+            if stochastic_skip:
+                self.skip_connection_merger = SkipConnectionMerger(
+                    channels=n_filters,
+                    nonlin=nonlin,
+                    batchnorm=batchnorm,
+                    dropout=dropout,
+                    res_block_type=res_block_type,
+                    merge_type=merge_type,
+                    conv2d_bias=conv2d_bias,
+                    res_block_kernel=res_block_kernel,
+                    res_block_skip_padding=res_block_skip_padding,
+                )
+
+        # print(f'[{self.__class__.__name__}] normalize_latent_factor:{self.normalize_latent_factor}')
+
+    def sample_from_q(
+        self,
+        input_: torch.Tensor,
+        bu_value: torch.Tensor,
+        var_clip_max: float = None,
+        mask: torch.Tensor = None,
+    ) -> torch.Tensor:
+        """
+        This method computes the latent inference distribution q(z_i|z_{i+1}) amd samples a latent tensor from it.
+
+        Parameters
+        ----------
+        input_: torch.Tensor
+            The input tensor to the layer, which is the output of the top-down layer above.
+        bu_value: torch.Tensor
+            The tensor defining the parameters /mu_q and /sigma_q computed during the bottom-up deterministic pass
+            at the correspondent hierarchical layer.
+        var_clip_max: float, optional
+            The maximum value reachable by the log-variance of the latent distribtion.
+            Values exceeding this threshold are clipped. Default is `None`.
+        mask: Union[None, torch.Tensor], optional
+            A tensor that is used to mask the sampled latent tensor. Default is `None`.
+        """
+        if self.is_top_layer:  # In top layer, we don't merge bu_value with p_params
+            q_params = bu_value
+        else:
+            # NOTE: Here the assumption is that the vampprior is only applied on the top layer.
+            n_img_prior = None
+            p_params = self.get_p_params(input_, n_img_prior)
+            q_params = self.merge(bu_value, p_params)
+
+        sample = self.stochastic.sample_from_q(q_params, var_clip_max)
+
+        if mask:
+            return sample[mask]
+
+        return sample
+
+    def get_p_params(
+        self,
+        input_: torch.Tensor,
+        n_img_prior: int,
+    ) -> torch.Tensor:
+        """
+        This method returns the parameters of the prior distribution p(z_i|z_{i+1}) for the latent tensor
+        depending on the hierarchical level of the layer and other specific conditions.
+
+        Parameters
+        ----------
+        input_: torch.Tensor
+            The input tensor to the layer, which is the output of the top-down layer above.
+        n_img_prior: int
+            The number of images to be generated from the unconditional prior distribution p(z_L).
+        """
+        p_params = None
+
+        # If top layer, define p_params as the ones of the prior p(z_L)
+        if self.is_top_layer:
+            p_params = self.top_prior_params
+
+            # Sample specific number of images by expanding the prior
+            if n_img_prior is not None:
+                p_params = p_params.expand(n_img_prior, -1, -1, -1)
+
+        # Else the input from the layer above is p_params itself
+        else:
+            p_params = input_
+
+        return p_params
+
+    def align_pparams_buvalue(
+        self, p_params: torch.Tensor, bu_value: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        In case the padding is not used either (or both) in encoder and decoder, we could have a shape mismatch
+        in the spatial dimensions (usually, dim=2 & dim=3).
+        This method performs a centercrop to ensure that both remain aligned.
+
+        Parameters
+        ----------
+        p_params: torch.Tensor
+            The tensor defining the parameters /mu_p and /sigma_p for the latent distribution p(z_i|z_{i+1}).
+        bu_value: torch.Tensor
+            The tensor defining the parameters /mu_q and /sigma_q computed during the bottom-up deterministic pass
+            at the correspondent hierarchical layer.
+        """
+        if bu_value.shape[-2:] != p_params.shape[-2:]:
+            assert self.bottomup_no_padding_mode is True
+            if self.topdown_no_padding_mode is False:
+                assert bu_value.shape[-1] > p_params.shape[-1]
+                bu_value = F.center_crop(bu_value, p_params.shape[-2:])
+            else:
+                if bu_value.shape[-1] > p_params.shape[-1]:
+                    bu_value = F.center_crop(bu_value, p_params.shape[-2:])
+                else:
+                    p_params = F.center_crop(p_params, bu_value.shape[-2:])
+        return p_params, bu_value
+
+    def forward(
+        self,
+        input_: torch.Tensor = None,
+        skip_connection_input: torch.Tensor = None,
+        inference_mode: bool = False,
+        bu_value: torch.Tensor = None,
+        n_img_prior: int = None,
+        forced_latent: torch.Tensor = None,
+        use_mode: bool = False,
+        force_constant_output: bool = False,
+        mode_pred: bool = False,
+        use_uncond_mode: bool = False,
+        var_clip_max: float = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor, Dict[str, torch.Tensor]]:
+        """
+        Parameters
+        ----------
+        input_: torch.Tensor, optional
+            The input tensor to the layer, which is the output of the top-down layer above.
+            Default is `None`.
+        skip_connection_input: torch.Tensor, optional
+            The tensor brought by the skip connection between the current and the previous top-down layer.
+            Default is `None`.
+        inference_mode: bool, optional
+            Whether the layer is in inference mode. See NOTE 2 in class description for more info.
+            Default is `False`.
+        bu_value: torch.Tensor, optional
+            The tensor defining the parameters /mu_q and /sigma_q computed during the bottom-up deterministic pass
+            at the correspondent hierarchical layer. Default is `None`.
+        n_img_prior: int, optional
+            The number of images to be generated from the unconditional prior distribution p(z_L).
+            Default is `None`.
+        forced_latent: torch.Tensor, optional
+            A pre-defined latent tensor. If it is not `None`, than it is used as the actual latent tensor and,
+            hence, sampling does not happen. Default is `None`.
+        use_mode: bool, optional
+            Wheteher the latent tensor should be set as the latent distribution mode.
+            In the case of Gaussian, the mode coincides with the mean of the distribution.
+            Default is `False`.
+        force_constant_output: bool, optional
+            Whether to copy the first sample (and rel. distrib parameters) over the whole batch.
+            This is used when doing experiment from the prior - q is not used.
+            Default is `False`.
+        mode_pred: bool, optional
+            Whether the model is in prediction mode. Default is `False`.
+        use_uncond_mode: bool, optional
+            Whether to use the uncoditional distribution p(z) to sample latents in prediction mode.
+        var_clip_max: float
+            The maximum value reachable by the log-variance of the latent distribtion.
+            Values exceeding this threshold are clipped.
+        """
+        # Check consistency of arguments
+        inputs_none = input_ is None and skip_connection_input is None
+        if self.is_top_layer and not inputs_none:
+            raise ValueError("In top layer, inputs should be None")
+
+        p_params = self.get_p_params(input_, n_img_prior)
+
+        # Get the parameters for the latent distribution to sample from
+        if inference_mode:
+            if self.is_top_layer:
+                q_params = bu_value
+                if mode_pred is False:
+                    p_params, bu_value = self.align_pparams_buvalue(p_params, bu_value)
+            else:
+                if use_uncond_mode:
+                    q_params = p_params
+                else:
+                    p_params, bu_value = self.align_pparams_buvalue(p_params, bu_value)
+                    q_params = self.merge(bu_value, p_params)
+        # In generative mode, q is not used
+        else:
+            q_params = None
+
+        # NOTE: Sampling is done either from q(z_i | z_{i+1}, x) or p(z_i | z_{i+1})
+        # depending on the mode (hence, in practice, by checking whether q_params is None).
+
+        # Normalization of latent space parameters:
+        # it is done, purely for stablity. See Very deep VAEs generalize autoregressive models.
+        if self.normalize_latent_factor:
+            q_params = q_params / self.normalize_latent_factor
+
+        # Sample (and process) a latent tensor in the stochastic layer
+        x, data_stoch = self.stochastic(
+            p_params=p_params,
+            q_params=q_params,
+            forced_latent=forced_latent,
+            use_mode=use_mode,
+            force_constant_output=force_constant_output,
+            analytical_kl=self.analytical_kl,
+            mode_pred=mode_pred,
+            use_uncond_mode=use_uncond_mode,
+            var_clip_max=var_clip_max,
+        )
+
+        # Merge skip connection from previous layer
+        if self.stochastic_skip and not self.is_top_layer:
+            if self.topdown_no_padding_mode is True:
+                # If no padding is done in the current top-down pass, there may be a shape mismatch between current tensor and skip connection input.
+                # As an example, if the output of last TopDownLayer was of size 64*64, due to lack of padding in the current layer, the current tensor
+                # might become different in shape, say 60*60.
+                # In order to avoid shape mismatch, we do central crop of the skip connection input.
+                skip_connection_input = F.center_crop(
+                    skip_connection_input, x.shape[-2:]
+                )
+
+            x = self.skip_connection_merger(x, skip_connection_input)
+
+        # Save activation before residual block as it can be the skip connection input in the next layer
+        x_pre_residual = x
+
+        if self.retain_spatial_dims:
+            # when we don't want to do padding in topdown as well, we need to spare some boundary pixels which would be used up.
+            extra_len = (self.topdown_no_padding_mode is True) * 3
+
+            # this means that x should be of the same size as config.data.image_size. So, we have to centercrop by a factor of 2 at this point.
+            # assert x.shape[-1] >= self.latent_shape[-1] // 2 + extra_len
+            # we assume that one topdown layer will have exactly one upscaling layer.
+            new_latent_shape = (
+                self.latent_shape[0] // 2 + extra_len,
+                self.latent_shape[1] // 2 + extra_len,
+            )
+
+            # If the LC is not applied on all layers, then this can happen.
+            if x.shape[-1] > new_latent_shape[-1]:
+                x = F.center_crop(x, new_latent_shape)
+
+        # Last top-down block (sequence of residual blocks)
+        x = self.deterministic_block(x)
+
+        if self.topdown_no_padding_mode:
+            x = F.center_crop(x, self.latent_shape)
+
+        # Save some metrics that will be used in the loss computation
+        keys = [
+            "z",
+            "kl_samplewise",
+            "kl_samplewise_restricted",
+            "kl_spatial",
+            "kl_channelwise",
+            # 'logprob_p',
+            "logprob_q",
+            "qvar_max",
+        ]
+        data = {k: data_stoch.get(k, None) for k in keys}
+        data["q_mu"] = None
+        data["q_lv"] = None
+        if data_stoch["q_params"] is not None:
+            q_mu, q_lv = data_stoch["q_params"]
+            data["q_mu"] = q_mu
+            data["q_lv"] = q_lv
+
+        return x, x_pre_residual, data
+
+
+class NormalStochasticBlock2d(nn.Module):
+    """
+    Stochastic block used in the Top-Down inference pass.
+
+    Algorithm:
+        - map input parameters to q(z) and (optionally) p(z) via convolution
+        - sample a latent tensor z ~ q(z)
+        - feed z to convolution and return.
+
+    NOTE 1:
+        If parameters for q are not given, sampling is done from p(z).
+
+    NOTE 2:
+        The restricted KL divergence is obtained by first computing the element-wise KL divergence
+        (i.e., the KL computed for each element of the latent tensors). Then, the restricted version
+        is computed by summing over the channels and the spatial dimensions associated only to the
+        portion of the latent tensor that is used for prediction.
+    """
+
+    def __init__(
+        self,
+        c_in: int,
+        c_vars: int,
+        c_out: int,
+        kernel: int = 3,
+        transform_p_params: bool = True,
+        vanilla_latent_hw: int = None,
+        restricted_kl: bool = False,
+        use_naive_exponential: bool = False,
+    ):
+        """
+        Parameters
+        ----------
+        c_in: int
+            The number of channels of the input tensor.
+        c_vars: int
+            The number of channels of the latent space tensor.
+        c_out:  int
+            The output of the stochastic layer.
+            Note that this is different from the sampled latent z.
+        kernel: int, optional
+            The size of the kernel used in convolutional layers.
+            Default is 3.
+        transform_p_params: bool, optional
+            Whether a transformation should be applied to the `p_params` tensor.
+            The transformation consists in a 2D convolution ()`conv_in_p()`) that
+            maps the input to a larger number of channels.
+            Default is `True`.
+        vanilla_latent_hw: int, optional
+            The shape of the latent tensor used for prediction (i.e., it influences the computation of restricted KL).
+            Default is `None`.
+        restricted_kl: bool, optional
+            Whether to compute the restricted version of KL Divergence.
+            See NOTE 2 for more information about its computation.
+            Default is `False`.
+        use_naive_exponential: bool, optional
+            If `False`, exponentials are computed according to the alternative definition
+            provided by `StableExponential` class. This should improve numerical stability
+            in the training process. Default is `False`.
+        """
+        super().__init__()
+        assert kernel % 2 == 1
+        pad = kernel // 2
+        self.transform_p_params = transform_p_params
+        self.c_in = c_in
+        self.c_out = c_out
+        self.c_vars = c_vars
+        self._use_naive_exponential = use_naive_exponential
+        self._vanilla_latent_hw = vanilla_latent_hw
+        self._restricted_kl = restricted_kl
+
+        if transform_p_params:
+            self.conv_in_p = nn.Conv2d(c_in, 2 * c_vars, kernel, padding=pad)
+        self.conv_in_q = nn.Conv2d(c_in, 2 * c_vars, kernel, padding=pad)
+        self.conv_out = nn.Conv2d(c_vars, c_out, kernel, padding=pad)
+
+    # def forward_swapped(self, p_params, q_mu, q_lv):
+    #
+    #     if self.transform_p_params:
+    #         p_params = self.conv_in_p(p_params)
+    #     else:
+    #         assert p_params.size(1) == 2 * self.c_vars
+    #
+    #     # Define p(z)
+    #     p_mu, p_lv = p_params.chunk(2, dim=1)
+    #     p = Normal(p_mu, (p_lv / 2).exp())
+    #
+    #     # Define q(z)
+    #     q = Normal(q_mu, (q_lv / 2).exp())
+    #     # Sample from q(z)
+    #     sampling_distrib = q
+    #
+    #     # Generate latent variable (typically by sampling)
+    #     z = sampling_distrib.rsample()
+    #
+    #     # Output of stochastic layer
+    #     out = self.conv_out(z)
+    #
+    #     data = {
+    #         'z': z,  # sampled variable at this layer (batch, ch, h, w)
+    #         'p_params': p_params,  # (b, ch, h, w) where b is 1 or batch size
+    #     }
+    #     return out, data
+
+    def get_z(
+        self,
+        sampling_distrib: torch.distributions.normal.Normal,
+        forced_latent: torch.Tensor,
+        use_mode: bool,
+        mode_pred: bool,
+        use_uncond_mode: bool,
+    ) -> torch.Tensor:
+        """
+        This method enables to sample a latent tensor given the distribution to sample from.
+
+        Latent variable can be obtained is several ways:
+            - Sampled from the (Gaussian) latent distribution.
+            - Taken as a pre-defined forced latent.
+            - Taken as the mode (mean) of the latent distribution.
+            - In prediction mode (`mode_pred==True`), can be either sample or taken as the distribution mode.
+
+        Parameters
+        ----------
+        sampling_distrib: torch.distributions.normal.Normal
+            The Gaussian distribution from which latent tensor is sampled.
+        forced_latent: torch.Tensor
+            A pre-defined latent tensor. If it is not `None`, than it is used as the actual latent tensor and,
+            hence, sampling does not happen.
+        use_mode: bool
+            Wheteher the latent tensor should be set as the latent distribution mode.
+            In the case of Gaussian, the mode coincides with the mean of the distribution.
+        mode_pred: bool
+            Whether the model is prediction mode.
+        use_uncond_mode: bool
+            Whether to use the uncoditional distribution p(z) to sample latents in prediction mode.
+        """
+        if forced_latent is None:
+            if use_mode:
+                z = sampling_distrib.mean
+            else:
+                if mode_pred:
+                    if use_uncond_mode:
+                        z = sampling_distrib.mean
+                    else:
+                        z = sampling_distrib.rsample()
+                else:
+                    z = sampling_distrib.rsample()
+        else:
+            z = forced_latent
+        return z
+
+    def sample_from_q(
+        self, q_params: torch.Tensor, var_clip_max: float
+    ) -> torch.Tensor:
+        """
+        Given an input parameter tensor defining q(z),
+        it processes it by calling `process_q_params()` method and
+        sample a latent tensor from the resulting distribution.
+
+        Parameters
+        ----------
+        q_params: torch.Tensor
+            The input tensor to be processed.
+        var_clip_max: float
+            The maximum value reachable by the log-variance of the latent distribtion.
+            Values exceeding this threshold are clipped.
+        """
+        _, _, q = self.process_q_params(q_params, var_clip_max)
+        return q.rsample()
+
+    def compute_kl_metrics(
+        self,
+        p: torch.distributions.normal.Normal,
+        p_params: torch.Tensor,
+        q: torch.distributions.normal.Normal,
+        q_params: torch.Tensor,
+        mode_pred: bool,
+        analytical_kl: bool,
+        z: torch.Tensor,
+    ) -> Dict[str, torch.Tensor]:
+        """
+        Compute KL (analytical or MC estimate) and then process it, extracting composed versions of the metric.
+        Specifically, the different versions of the KL loss terms are:
+            - `kl_elementwise`: KL term for each single element of the latent tensor [Shape: (batch, ch, h, w)].
+            - `kl_samplewise`: KL term associated to each sample in the batch [Shape: (batch, )].
+            - `kl_samplewise_restricted`: KL term only associated to the portion of the latent tensor that is
+            used for prediction and summed over channel and spatial dimensions [Shape: (batch, )].
+            - `kl_channelwise`: KL term associated to each sample and each channel [Shape: (batch, ch, )].
+            - `kl_spatial`: KL term summed over the channels, i.e., retaining the spatial dimensions [Shape: (batch, h, w)]
+
+        Parameters
+        ----------
+        p: torch.distributions.normal.Normal
+            The prior generative distribution p(z_i|z_{i+1}) (or p(z_L)).
+        p_params: torch.Tensor
+            The parameters of the prior generative distribution.
+        q: torch.distributions.normal.Normal
+            The inference distribution q(z_i|z_{i+1}) (or q(z_L|x)).
+        q_params: torch.Tensor
+            The parameters of the inference distribution.
+        mode_pred: bool
+            Whether the model is in prediction mode.
+        analytical_kl: bool
+            Whether to compute the KL divergence analytically or using Monte Carlo estimation.
+        z: torch.Tensor
+            The sampled latent tensor.
+        """
+        kl_samplewise_restricted = None
+
+        if mode_pred is False:  # if not in prediction mode
+            # KL term for each single element of the latent tensor [Shape: (batch, ch, h, w)]
+            if analytical_kl:
+                kl_elementwise = kl_divergence(q, p)
+            else:
+                kl_elementwise = kl_normal_mc(z, p_params, q_params)
+
+            # KL term only associated to the portion of the latent tensor that is used for prediction and
+            # summed over channel and spatial dimensions. [Shape: (batch, )]
+            # NOTE: vanilla_latent_hw is the shape of the latent tensor used for prediction, hence
+            # the restriction has shape [Shape: (batch, ch, vanilla_latent_hw[0], vanilla_latent_hw[1])]
+            if self._restricted_kl:
+                pad = (kl_elementwise.shape[-1] - self._vanilla_latent_hw) // 2
+                assert pad > 0, "Disable restricted kl since there is no restriction."
+                tmp = kl_elementwise[..., pad:-pad, pad:-pad]
+                kl_samplewise_restricted = tmp.sum((1, 2, 3))
+
+            # KL term associated to each sample in the batch [Shape: (batch, )]
+            kl_samplewise = kl_elementwise.sum((1, 2, 3))
+
+            # KL term associated to each sample and each channel [Shape: (batch, ch, )]
+            kl_channelwise = kl_elementwise.sum((2, 3))
+
+            # KL term summed over the channels, i.e., retaining the spatial dimensions [Shape: (batch, h, w)]
+            kl_spatial = kl_elementwise.sum(1)
+        else:  # if predicting, no need to compute KL
+            kl_elementwise = kl_samplewise = kl_spatial = kl_channelwise = None
+
+        kl_dict = {
+            "kl_elementwise": kl_elementwise,  # (batch, ch, h, w)
+            "kl_samplewise": kl_samplewise,  # (batch, )
+            "kl_samplewise_restricted": kl_samplewise_restricted,  # (batch, )
+            "kl_channelwise": kl_channelwise,  # (batch, ch)
+            "kl_spatial": kl_spatial,  # (batch, h, w)
+        }
+        return kl_dict
+
+    def process_p_params(
+        self, p_params: torch.Tensor, var_clip_max: float
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.distributions.normal.Normal]:
+        """
+        Process the input parameters to get the prior distribution p(z_i|z_{i+1}) (or p(z_L)).
+
+        Processing consists in:
+            - (optionally) 2D convolution on the input tensor to increase number of channels.
+            - split the resulting tensor into two chunks, the mean and the log-variance.
+            - (optionally) clip the log-variance to an upper threshold.
+            - define the normal distribution p(z) given the parameter tensors above.
+
+        Parameters
+        ----------
+        p_params: torch.Tensor
+            The input tensor to be processed.
+        var_clip_max: float
+            The maximum value reachable by the log-variance of the latent distribtion.
+            Values exceeding this threshold are clipped.
+        """
+        if self.transform_p_params:
+            p_params = self.conv_in_p(p_params)
+        else:
+            assert p_params.size(1) == 2 * self.c_vars
+
+        # Define p(z)
+        p_mu, p_lv = p_params.chunk(2, dim=1)
+        if var_clip_max is not None:
+            p_lv = torch.clip(p_lv, max=var_clip_max)
+
+        p_mu = StableMean(p_mu)
+        p_lv = StableLogVar(p_lv, enable_stable=not self._use_naive_exponential)
+        p = Normal(p_mu.get(), p_lv.get_std())
+        return p_mu, p_lv, p
+
+    def process_q_params(
+        self, q_params: torch.Tensor, var_clip_max: float, allow_oddsizes: bool = False
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.distributions.normal.Normal]:
+        """
+        Process the input parameters to get the inference distribution q(z_i|z_{i+1}) (or q(z|x)).
+
+        Processing consists in:
+            - 2D convolution on the input tensor to increase number of channels.
+            - split the resulting tensor into two chunks, the mean and the log-variance.
+            - (optionally) clip the log-variance to an upper threshold.
+            - (optionally) crop the resulting tensors to ensure that the last spatial dimension is even.
+            - define the normal distribution q(z) given the parameter tensors above.
+
+        Parameters
+        ----------
+        p_params: torch.Tensor
+            The input tensor to be processed.
+        var_clip_max: float
+            The maximum value reachable by the log-variance of the latent distribtion.
+            Values exceeding this threshold are clipped.
+        """
+        q_params = self.conv_in_q(q_params)
+
+        q_mu, q_lv = q_params.chunk(2, dim=1)
+        if var_clip_max is not None:
+            q_lv = torch.clip(q_lv, max=var_clip_max)
+
+        if q_mu.shape[-1] % 2 == 1 and allow_oddsizes is False:
+            q_mu = F.center_crop(q_mu, q_mu.shape[-1] - 1)
+            q_lv = F.center_crop(q_lv, q_lv.shape[-1] - 1)
+            # clip_start = np.random.rand() > 0.5
+            # q_mu = q_mu[:, :, 1:, 1:] if clip_start else q_mu[:, :, :-1, :-1]
+            # q_lv = q_lv[:, :, 1:, 1:] if clip_start else q_lv[:, :, :-1, :-1]
+
+        q_mu = StableMean(q_mu)
+        q_lv = StableLogVar(q_lv, enable_stable=not self._use_naive_exponential)
+        q = Normal(q_mu.get(), q_lv.get_std())
+        return q_mu, q_lv, q
+
+    def forward(
+        self,
+        p_params: torch.Tensor,
+        q_params: torch.Tensor = None,
+        forced_latent: torch.Tensor = None,
+        use_mode: bool = False,
+        force_constant_output: bool = False,
+        analytical_kl: bool = False,
+        mode_pred: bool = False,
+        use_uncond_mode: bool = False,
+        var_clip_max: float = None,
+    ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
+        """
+        Parameters
+        ----------
+        p_params: torch.Tensor
+            The output tensor of the top-down layer above (i.e., mu_{p,i+1}, sigma_{p,i+1}).
+        q_params: torch.Tensor, optional
+            The tensor resulting from merging the bu_value tensor at the same hierarchical level
+            from the bottom-up pass and the `p_params` tensor. Default is `None`.
+        forced_latent: torch.Tensor, optional
+            A pre-defined latent tensor. If it is not `None`, than it is used as the actual latent
+            tensor and, hence, sampling does not happen. Default is `None`.
+        use_mode: bool, optional
+            Wheteher the latent tensor should be set as the latent distribution mode.
+            In the case of Gaussian, the mode coincides with the mean of the distribution.
+            Default is `False`.
+        force_constant_output: bool, optional
+            Whether to copy the first sample (and rel. distrib parameters) over the whole batch.
+            This is used when doing experiment from the prior - q is not used.
+            Default is `False`.
+        analytical_kl: bool, optional
+            Whether to compute the KL divergence analytically or using Monte Carlo estimation.
+            Default is `False`.
+        mode_pred: bool, optional
+            Whether the model is in prediction mode. Default is `False`.
+        use_uncond_mode: bool, optional
+            Whether to use the uncoditional distribution p(z) to sample latents in prediction mode.
+            Default is `False`.
+        var_clip_max: float, optional
+            The maximum value reachable by the log-variance of the latent distribtion.
+            Values exceeding this threshold are clipped. Default is `None`.
+        """
+        debug_qvar_max = 0
+
+        # Check sampling options consistency
+        assert (forced_latent is None) or (not use_mode)
+
+        # Get generative distribution p(z_i|z_{i+1})
+        p_mu, p_lv, p = self.process_p_params(p_params, var_clip_max)
+        p_params = (p_mu, p_lv)
+
+        if q_params is not None:
+            # Get inference distribution q(z_i|z_{i+1})
+            # NOTE: At inference time, don't centercrop the q_params even if they are odd in size.
+            q_mu, q_lv, q = self.process_q_params(
+                q_params, var_clip_max, allow_oddsizes=mode_pred is True
+            )
+            q_params = (q_mu, q_lv)
+            sampling_distrib = q
+            debug_qvar_max = torch.max(q_lv.get())
+
+            # Centercrop p_params so that their size matches the one of q_params
+            q_size = q_mu.get().shape[-1]
+            if p_mu.get().shape[-1] != q_size and mode_pred is False:
+                p_mu.centercrop_to_size(q_size)
+                p_lv.centercrop_to_size(q_size)
+        else:
+            sampling_distrib = p
+
+        # Sample latent variable
+        z = self.get_z(
+            sampling_distrib, forced_latent, use_mode, mode_pred, use_uncond_mode
+        )
+
+        # Copy one sample (and distrib parameters) over the whole batch.
+        # This is used when doing experiment from the prior - q is not used.
+        if force_constant_output:
+            z = z[0:1].expand_as(z).clone()
+            p_params = (
+                p_params[0][0:1].expand_as(p_params[0]).clone(),
+                p_params[1][0:1].expand_as(p_params[1]).clone(),
+            )
+
+        # Pass the sampled latent througn the output convolutional layer of stochastic block
+        out = self.conv_out(z)
+
+        # Compute log p(z)# NOTE: disabling its computation.
+        # if mode_pred is False:
+        #     logprob_p =  p.log_prob(z).sum((1, 2, 3))
+        # else:
+        #     logprob_p = None
+
+        if q_params is not None:
+            # Compute log q(z)
+            logprob_q = q.log_prob(z).sum((1, 2, 3))
+            # Compute KL divergence metrics
+            kl_dict = self.compute_kl_metrics(
+                p, p_params, q, q_params, mode_pred, analytical_kl, z
+            )
+        else:
+            kl_dict = {}
+            logprob_q = None
+
+        # Store meaningful quantities to use them in following layers
+        data = kl_dict
+        data["z"] = z  # sampled variable at this layer (batch, ch, h, w)
+        data["p_params"] = p_params  # (b, ch, h, w) where b is 1 or batch size
+        data["q_params"] = q_params  # (batch, ch, h, w)
+        # data['logprob_p'] = logprob_p  # (batch, )
+        data["logprob_q"] = logprob_q  # (batch, )
+        data["qvar_max"] = debug_qvar_max
+
+        return out, data
+
+
+class NonStochasticBlock2d(nn.Module):
+    """
+    Non-stochastic version of the NormalStochasticBlock2d.
+    """
+
+    def __init__(
+        self,
+        c_vars: int,
+        c_in: int,
+        c_out: int,
+        kernel: int = 3,
+        groups: int = 1,
+        conv2d_bias: bool = True,
+        transform_p_params: bool = True,
+    ):
+        """
+        Constructor.
+
+        Parameters
+        ----------
+        c_vars: int
+            The number of channels of the latent space tensor.
+        c_in: int
+            The number of channels of the input tensor.
+        c_out:  int
+            The output of the stochastic layer.
+            Note that this is different from the sampled latent z.
+        kernel: int, optional
+            The size of the kernel used in convolutional layers.
+            Default is 3.
+        groups: int, optional
+            The number of groups to consider in the convolutions of this layer.
+            Default is 1.
+        conv2d_bias: bool, optional
+            Whether to use bias term is the convolutional blocks of this layer.
+            Default is `True`.
+        transform_p_params: bool, optional
+            Whether a transformation should be applied to the `p_params` tensor.
+            The transformation consists in a 2D convolution ()`conv_in_p()`) that
+            maps the input to a larger number of channels.
+            Default is `True`.
+        """
+        super().__init__()
+        assert kernel % 2 == 1
+        pad = kernel // 2
+        self.transform_p_params = transform_p_params
+        self.c_in = c_in
+        self.c_out = c_out
+        self.c_vars = c_vars
+
+        if transform_p_params:
+            self.conv_in_p = nn.Conv2d(
+                c_in, 2 * c_vars, kernel, padding=pad, bias=conv2d_bias, groups=groups
+            )
+        self.conv_in_q = nn.Conv2d(
+            c_in, 2 * c_vars, kernel, padding=pad, bias=conv2d_bias, groups=groups
+        )
+        self.conv_out = nn.Conv2d(
+            c_vars, c_out, kernel, padding=pad, bias=conv2d_bias, groups=groups
+        )
+
+    def compute_kl_metrics(
+        self,
+        p: torch.distributions.normal.Normal,
+        p_params: torch.Tensor,
+        q: torch.distributions.normal.Normal,
+        q_params: torch.Tensor,
+        mode_pred: bool,
+        analytical_kl: bool,
+        z: torch.Tensor,
+    ) -> Dict[str, None]:
+        """
+        Compute KL (analytical or MC estimate) and then process it, extracting composed versions of the metric.
+        Specifically, the different versions of the KL loss terms are:
+            - `kl_elementwise`: KL term for each single element of the latent tensor [Shape: (batch, ch, h, w)].
+            - `kl_samplewise`: KL term associated to each sample in the batch [Shape: (batch, )].
+            - `kl_samplewise_restricted`: KL term only associated to the portion of the latent tensor that is
+            used for prediction and summed over channel and spatial dimensions [Shape: (batch, )].
+            - `kl_channelwise`: KL term associated to each sample and each channel [Shape: (batch, ch, )].
+            - `kl_spatial`: # KL term summed over the channels, i.e., retaining the spatial dimensions [Shape: (batch, h, w)]
+
+        NOTE: in this class all the KL metrics are set to `None`.
+
+        Parameters
+        ----------
+        p: torch.distributions.normal.Normal
+            The prior generative distribution p(z_i|z_{i+1}) (or p(z_L)).
+        p_params: torch.Tensor
+            The parameters of the prior generative distribution.
+        q: torch.distributions.normal.Normal
+            The inference distribution q(z_i|z_{i+1}) (or q(z_L|x)).
+        q_params: torch.Tensor
+            The parameters of the inference distribution.
+        mode_pred: bool
+            Whether the model is in prediction mode.
+        analytical_kl: bool
+            Whether to compute the KL divergence analytically or using Monte Carlo estimation.
+        z: torch.Tensor
+            The sampled latent tensor.
+        """
+        kl_dict = {
+            "kl_elementwise": None,  # (batch, ch, h, w)
+            "kl_samplewise": None,  # (batch, )
+            "kl_spatial": None,  # (batch, h, w)
+            "kl_channelwise": None,  # (batch, ch)
+        }
+        return kl_dict
+
+    def process_p_params(self, p_params, var_clip_max):
+        if self.transform_p_params:
+            p_params = self.conv_in_p(p_params)
+        else:
+
+            assert (
+                p_params.size(1) == 2 * self.c_vars
+            ), f"{p_params.shape} {self.c_vars}"
+
+        # Define p(z)
+        p_mu, p_lv = p_params.chunk(2, dim=1)
+        return p_mu, None
+
+    def process_q_params(self, q_params, var_clip_max, allow_oddsizes=False):
+        # Define q(z)
+        q_params = self.conv_in_q(q_params)
+        q_mu, q_lv = q_params.chunk(2, dim=1)
+
+        if q_mu.shape[-1] % 2 == 1 and allow_oddsizes is False:
+            q_mu = F.center_crop(q_mu, q_mu.shape[-1] - 1)
+
+        return q_mu, None
+
+    def forward(
+        self,
+        p_params: torch.Tensor,
+        q_params: torch.Tensor = None,
+        forced_latent: Union[None, torch.Tensor] = None,
+        use_mode: bool = False,
+        force_constant_output: bool = False,
+        analytical_kl: bool = False,
+        mode_pred: bool = False,
+        use_uncond_mode: bool = False,
+        var_clip_max: float = None,
+    ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
+        """
+        Parameters
+        ----------
+        p_params: torch.Tensor
+            The output tensor of the top-down layer above (i.e., mu_{p,i+1}, sigma_{p,i+1}).
+        q_params: torch.Tensor, optional
+            The tensor resulting from merging the bu_value tensor at the same hierarchical level
+            from the bottom-up pass and the `p_params` tensor. Default is `None`.
+        forced_latent: torch.Tensor, optional
+            A pre-defined latent tensor. If it is not `None`, than it is used as the actual latent
+            tensor and, hence, sampling does not happen. Default is `None`.
+        use_mode: bool, optional
+            Wheteher the latent tensor should be set as the latent distribution mode.
+            In the case of Gaussian, the mode coincides with the mean of the distribution.
+            Default is `False`.
+        force_constant_output: bool, optional
+            Whether to copy the first sample (and rel. distrib parameters) over the whole batch.
+            This is used when doing experiment from the prior - q is not used.
+            Default is `False`.
+        analytical_kl: bool, optional
+            Whether to compute the KL divergence analytically or using Monte Carlo estimation.
+            Default is `False`.
+        mode_pred: bool, optional
+            Whether the model is in prediction mode. Default is `False`.
+        use_uncond_mode: bool, optional
+            Whether to use the uncoditional distribution p(z) to sample latents in prediction mode.
+            Default is `False`.
+        var_clip_max: float, optional
+            The maximum value reachable by the log-variance of the latent distribtion.
+            Values exceeding this threshold are clipped. Default is `None`.
+        """
+        debug_qvar_max = 0
+        assert (forced_latent is None) or (not use_mode)
+
+        p_mu, _ = self.process_p_params(p_params, var_clip_max)
+
+        p_params = (p_mu, None)
+
+        if q_params is not None:
+            # At inference time, just don't centercrop the q_params even if they are odd in size.
+            q_mu, _ = self.process_q_params(
+                q_params, var_clip_max, allow_oddsizes=mode_pred is True
+            )
+            q_params = (q_mu, None)
+            debug_qvar_max = torch.Tensor([1]).to(q_mu.device)
+            # Sample from q(z)
+            sampling_distrib = q_mu
+            q_size = q_mu.shape[-1]
+            if p_mu.shape[-1] != q_size and mode_pred is False:
+                p_mu.centercrop_to_size(q_size)
+        else:
+            # Sample from p(z)
+            sampling_distrib = p_mu
+
+        # Generate latent variable (typically by sampling)
+        z = sampling_distrib
+
+        # Copy one sample (and distrib parameters) over the whole batch.
+        # This is used when doing experiment from the prior - q is not used.
+        if force_constant_output:
+            z = z[0:1].expand_as(z).clone()
+            p_params = (
+                p_params[0][0:1].expand_as(p_params[0]).clone(),
+                p_params[1][0:1].expand_as(p_params[1]).clone(),
+            )
+
+        # Output of stochastic layer
+        out = self.conv_out(z)
+
+        kl_dict = {}
+        logprob_q = None
+
+        data = kl_dict
+        data["z"] = z  # sampled variable at this layer (batch, ch, h, w)
+        data["p_params"] = p_params  # (b, ch, h, w) where b is 1 or batch size
+        data["q_params"] = q_params  # (batch, ch, h, w)
+        data["logprob_q"] = logprob_q  # (batch, )
+        data["qvar_max"] = debug_qvar_max
+
+        return out, data