Hierarchical recurrent reasoning in one compact PyTorch file: two transformer modules updating at different rates, one-step-gradient recurrence, deep supervision and a Q-based halting head, with encoders for Sudoku, mazes and ARC grids. It follows the published Hierarchical Reasoning Model.
The whole model is one 31 KB file, and its dependencies are torch, numpy and tqdm. Inside it: a high-level and a low-level transformer module updating at different rates, rotary positional encoding, RMSNorm, one-step-gradient recurrence, deep supervision across segments, and a Q-based halting head that decides when to stop thinking. Encoders and decoders cover Sudoku, mazes and ARC grids, and the CLI is four verbs — train, solve, analyze, benchmark. The defaults live in code: 512 hidden, eight heads, four layers, a 1,024 vocabulary, two cycles of four steps, at most eight segments.
The engineering is defensive. Dataset loading fails closed: a missing or malformed path raises instead of silently substituting random data, which is the failure that produces a model that trains and means nothing. Spatial augmentation applies jointly to input and target, so a rotated puzzle keeps its answer. Attention runs on PyTorch's scaled dot-product path, with an explicit reference implementation kept for numerical parity and checked against it in CI.
The metrics are precise about what they mean. The benchmark reports cell accuracy and exact match; Sudoku adds a constraint-valid rate that requires every row, column and box to hold one to nine exactly once with the original clues unchanged. ARC is exact grid match. A change to recurrence, halting, augmentation or a metric arrives with a failing test first.