import torch


def flatten_tensors(x, len_shape):
    """Flatten tensor / list of tensors into a single tensor

    Args
        x
            tensor [*shape, *dims]

            or

            list of N tensors
            [*shape, *dims_1]
            ...
            [*shape, *dims_N]
        len_shape (int) is the length of shape

    Returns
        [*shape, total_ndims]

        OR

        [*shape, total_ndims_1 + ... + total_ndims_N]
    """
    # Flatten
    # --Flatten x
    if torch.is_tensor(x):
        shape = x.shape[:len_shape]

        # [*shape, total_ndims]
        x_flattened = x.view(*[*shape, -1])
    else:
        shape = x[0].shape[:len_shape]

        # [*shape, total_ndims_1 + ... + total_ndims_N]
        x_flattened = torch.cat([single_x.view(*[*shape, -1]) for single_x in x], dim=-1)

    return x_flattened