import math
import torch


def subbatch(toks, maxlen):
    """Splits tokens into sub-batches each containing at most maxlen tokens."""
    assert(maxlen > 0)
    _, dlen = toks.shape[:2]
    subbatch_qty = math.ceil(dlen / maxlen)
    S = math.ceil(dlen / subbatch_qty) if subbatch_qty > 0 else 0 # minimize the size given the number of subbatches
    stack = []
    if subbatch_qty == 1:
        return toks, subbatch_qty
    else:
        for s in range(subbatch_qty):
            stack.append(toks[:, s*S:(s+1)*S])
            if stack[-1].shape[1] != S:
                nulls = torch.zeros_like(toks[:, :S - stack[-1].shape[1]])
                stack[-1] = torch.cat([stack[-1], nulls], dim=1)
        return torch.cat(stack, dim=0), subbatch_qty