import torch


def finalize(s0, s1, s2):
    """Concatenate scattering of different orders.

    Parameters
    ----------
    s0 : tensor
        Tensor which contains the zeroth order scattering coefficents.
    s1 : tensor
        Tensor which contains the first order scattering coefficents.
    s2 : tensor
        Tensor which contains the second order scattering coefficents.
    
    Returns
    -------
    s : tensor
        Final output. Scattering transform.

    """
    if len(s2)>0:
        return torch.cat([torch.cat(s0, -2), torch.cat(s1, -2), torch.cat(s2, -2)], -2)
    else:
        return torch.cat([torch.cat(s0, -2), torch.cat(s1, -2)], -2)