tasks.py

"""Self-contained synthetic tasks. No data download needed.

Each task presents a sequence of symbols, a separator, then a block of MASK
slots. The target at the masked slots is some function of the input (copy,
reverse, or sort). The model must fill all masked slots in parallel, which
exercises both the associative routing of the equilibrium core and the
non-autoregressive decoder.
"""

import torch

PAD, MASK, SEP = 0, 1, 2
OFFSET = 3  # data symbols start here


def vocab_size(n_symbols: int) -> int:
    return OFFSET + n_symbols


def make_batch(task: str, batch: int, length: int, n_symbols: int, device):
    sym = torch.randint(OFFSET, OFFSET + n_symbols, (batch, length), device=device)
    if task == "copy":
        tgt_vals = sym
    elif task == "reverse":
        tgt_vals = torch.flip(sym, dims=[1])
    elif task == "sort":
        tgt_vals, _ = torch.sort(sym, dim=1)
    else:
        raise ValueError(f"unknown task: {task}")

    sep = torch.full((batch, 1), SEP, device=device, dtype=torch.long)
    masked = torch.full_like(tgt_vals, MASK)
    inp = torch.cat([sym, sep, masked], dim=1)
    tgt = torch.cat([sym, sep, tgt_vals], dim=1)

    loss_mask = torch.zeros_like(inp, dtype=torch.bool)
    loss_mask[:, length + 1:] = True  # loss only on the output block
    return inp, tgt, loss_mask