def split_into_chunks(batch, chunk_sizes):
  """Returns batch split in 3 according to chunk_sizes.

  Args:
    batch: A sequence of length sum(chunk_sizes), usually examples or targets.
    chunk_sizes: A tuple of 3 ints (flush_size, support_size, query_size).

  Returns:
    A tuple of 3 sequences (flush_chunk, support_chunk, query_chunk).
  """
  assert sum(chunk_sizes) == len(batch)
  flush_chunk_size, support_chunk_size, _ = chunk_sizes
  query_start = flush_chunk_size + support_chunk_size

  flush_chunk = batch[:flush_chunk_size]
  support_chunk = batch[flush_chunk_size:query_start]
  query_chunk = batch[query_start:]

  return (flush_chunk, support_chunk, query_chunk)