import torch


def assemble_recon_data(heldin, heldin_forward, heldout, heldout_forward):
    """Combines heldin/heldout and observed/forward data
    into a single tensor.

    Parameters
    ----------
    heldin : torch.Tensor
        A BxTxN tensor.
    heldin_forward : torch.Tensor
        A BxT_FWDxN tensor.
    heldout : torch.Tensor
        A BxTxN_OUT tensor.
    heldout_forward : torch.Tensor
        A BxT_FWDxN_OUT tensor.

    Returns
    -------
    torch.tensor
        A Bx(T+T_FWD)x(N+N_OUT) tensor.
    """
    heldin_full = torch.cat([heldin, heldin_forward], dim=1)
    heldout_full = torch.cat([heldout, heldout_forward], dim=1)
    recon_data = torch.cat([heldin_full, heldout_full], dim=2)
    return recon_data