import torch


def divide_sequence(
        x: torch.Tensor,
        seq_len: int,
        pad: bool) -> torch.Tensor:
    """
    Break full_len -> n_samples * seq_len

    Args:
        x:
            tensor: (full_len, ...)

        seq_len:
            Divided sequence length, the second dimension of the output tensor

        pad:
            Pad with zeros or discard the remainder sequence

    Returns:
        tensor, where the first input dimension (full_len, ) is split into (n_samples, seq_len)
    """

    full_len = x.size()[0]
    k_dims = list(x.size()[1:])

    remainder = full_len % seq_len
    divisible = remainder == 0

    if not divisible:
        if pad:
            pad_len = seq_len - remainder

            pad_size = [pad_len] + k_dims
            x_pad = torch.zeros(size=pad_size, dtype=x.dtype)

            if x.is_cuda:
                x_pad = x_pad.cuda()

            x = torch.cat([x, x_pad], dim=0)

        else:  # discard remainder
            x = x[0:-remainder]

    new_size = [-1, seq_len] + k_dims

    return x.view(*new_size)