def get_batch(chunk_agg: list,  # chunk aggregate from get_chunks function
              batch_size: int,  # how many samples to load
              i: int,  # current batch index
              return_malicious: bool = False,  # whether to return the malicious label for the data points or not
              return_counts: bool = False,  # whether to return the counts for the data points or not
              return_tags: bool = False):  # whether to return the tags for the data points or not
    """ Get a batch of data from a chunk aggregate.

    Args:
        chunk_agg: Chunk aggregate from get_chunks function
        batch_size: How many samples to load
        i: Current batch index
        return_malicious: Whether to return the malicious label for the data points or not (default: False)
        return_counts: Whether to return the counts for the data points or not (default: False)
        return_tags: Whether to return the tags for the data points or not (default: False)
    Returns:
        Current batch of sha (optional), features and labels.
    """

    # get current batch of data using i and batch size
    batch = [t[i:i + batch_size] for t in chunk_agg]
    # pop the last element of the current batch (y -> labels)
    batch_y = batch.pop()
    # initialize labels dict
    labels = {}

    if return_malicious:
        # get malware label for this sample through the index
        labels['malware'] = batch_y[:, 0]

    if return_counts:
        # get count for this sample through the index
        labels['count'] = batch_y[:, 1]

    if return_tags:
        # get tags list for this sample through the index
        labels['tags'] = batch_y[:, 2:]

    # return current batch unpacked (contains S (optionally) and X) and labels dict
    return *batch, labels