def split_xz(xz, x_dims, x_only=False, z_only=False):
    """
    Split concatenated xz vector into x and z vectors.

    Args:
        xz (list): The XZ matrix.
        x_dims ([list/tuple]) the dimensions of the X dimensions
        x_only (bool): If True, returns only the x vector.
        z_only (bool): If True, returns only the z vector.

    Returns:
        x, z (list, list): the separate X and Z matrices.

    """
    if x_only and z_only:
        raise ValueError(
            "Please select either x_only or z_only or set both"
            "to false to return both x and z after splitting."
        )
    x, z = xz[: len(x_dims)], xz[len(x_dims) :]
    if x_only:
        return x
    elif z_only:
        return z
    else:
        return x, z