This is an implementation of the Tiny Recursive Model (TRM)

Reference code

ARC-AGI: You need to download the data from the kaggle challenge page.

Other datasets are generated locally.

To generate a Sudoku dataset: ./bash/generate_sudoku_data.sh

uv run python src/nn/train.py experiment=trm_sudoku4x4 (takes a few mins on an A10)

uv run python src/nn/evaluate.py +checkpoint=/tmp/ml-experiments/lunar-pine-174/checkpoints/last.ckpt +data_dir=./data/sudoku_4x4_small

To visualize the results of a model, use notebooks/neural_viewer.ipynb

Follow me on X