{"id":32933395,"url":"https://github.com/allthingssecurity/trm_sudoku","last_synced_at":"2026-05-17T04:31:50.420Z","repository":{"id":323501249,"uuid":"1093506562","full_name":"allthingssecurity/trm_sudoku","owner":"allthingssecurity","description":"TRM Sudoku 4x4 – Tiny Recursive Model that solves Sudoku on macOS (MPS) with PyTorch Lightning and Hydra.","archived":false,"fork":false,"pushed_at":"2025-11-10T13:31:05.000Z","size":224,"stargazers_count":0,"open_issues_count":0,"forks_count":0,"subscribers_count":0,"default_branch":"main","last_synced_at":"2025-11-10T15:18:59.949Z","etag":null,"topics":["deep-learning","hydra","lightning","mps","pytorch","sudoku","tiny-recursive-model","transformer","trm"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/allthingssecurity.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null,"zenodo":null,"notice":null,"maintainers":null,"copyright":null,"agents":null,"dco":null,"cla":null}},"created_at":"2025-11-10T13:18:37.000Z","updated_at":"2025-11-10T13:31:09.000Z","dependencies_parsed_at":null,"dependency_job_id":null,"html_url":"https://github.com/allthingssecurity/trm_sudoku","commit_stats":null,"previous_names":["allthingssecurity/trm_sudoku"],"tags_count":null,"template":false,"template_full_name":null,"purl":"pkg:github/allthingssecurity/trm_sudoku","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/allthingssecurity%2Ftrm_sudoku","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/allthingssecurity%2Ftrm_sudoku/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/allthingssecurity%2Ftrm_sudoku/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/allthingssecurity%2Ftrm_sudoku/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/allthingssecurity","download_url":"https://codeload.github.com/allthingssecurity/trm_sudoku/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/allthingssecurity%2Ftrm_sudoku/sbom","scorecard":null,"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":283910072,"owners_count":26915128,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","status":"online","status_checked_at":"2025-11-11T02:00:06.610Z","response_time":65,"last_error":null,"robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":true,"can_crawl_api":true,"host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["deep-learning","hydra","lightning","mps","pytorch","sudoku","tiny-recursive-model","transformer","trm"],"created_at":"2025-11-11T19:00:53.726Z","updated_at":"2025-11-11T19:01:58.505Z","avatar_url":"https://github.com/allthingssecurity.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"# An implementation of TRM\n\nThis is an implementation of the [Tiny Recursive Model (TRM)](https://arxiv.org/pdf/2510.04871v1)\n\nReference [code](https://github.com/SamsungSAILMontreal/TinyRecursiveModels)\n\n# Scope\n\nThis repo focuses on training and solving Sudoku 4x4 with TRM. No external datasets are required — puzzles are generated on the fly.\n\n# Setup\n\nOn macOS, install `uv` via Homebrew or the official installer:\n\n- Homebrew: `brew install uv`\n- Script: `curl -LsSf https://astral.sh/uv/install.sh | sh`\n\nAlternatively you can use a virtual environment with `python -m venv .venv \u0026\u0026 source .venv/bin/activate` and `pip install -e .`.\n\n# Training\n\nQuick GPU sanity run (macOS MPS):\n\n```bash\nuv run python src/nn/train.py \\\n  experiment=trm_sudoku4x4 \\\n  trainer=gpu trainer.accelerator=mps \\\n  timekeeping.max_epochs=5\n```\n\nStable longer run (good results):\n\n```bash\nuv run python src/nn/train.py \\\n  experiment=trm_sudoku4x4 \\\n  trainer=gpu trainer.accelerator=mps trainer.precision=32-true \\\n  timekeeping.max_epochs=60 timekeeping.batch_size=128 \\\n  model_tuning.hidden_size=128 model_tuning.num_layers=2 \\\n  model_tuning.N_supervision=2 \\\n  model_tuning.learning_rate=3e-4 model_tuning.learning_rate_emb=3e-3\n```\n\n# Evaluation\n\nSudoku is evaluated via the validation loop and simple scripts. Example to validate a checkpoint:\n\n```bash\nuv run python - \u003c\u003c'PY'\nfrom lightning import Trainer\nimport torch\nfrom src.nn.data.sudoku4x4_datamodule import SudokuDataModule\nfrom src.nn.models.trm_module import TRMModule\n\nckpt = 'train/runs/\u003ctimestamp\u003e/checkpoints/last.ckpt'\n\ndm = SudokuDataModule(\n    batch_size=128, num_workers=0,\n    grid_size=4, max_grid_size=6,\n    generate_on_fly=True,\n    num_train_puzzles=2000, num_val_puzzles=800,\n)\ndm.setup('fit')\n\ndevice = 'mps' if torch.backends.mps.is_available() else 'cpu'\nmodel = TRMModule.load_from_checkpoint(ckpt, map_location=device)\n\nprint(Trainer(accelerator=device, devices=1).validate(model, dm.val_dataloader()))\nPY\n```\n\n# Sample Predictions\n\nHere is a short script to print a few held-out Sudoku4x4 puzzles with predictions:\n\n```bash\nuv run python - \u003c\u003c'PY'\nimport torch, numpy as np\nfrom src.nn.data.sudoku4x4_datamodule import SudokuDataModule\nfrom src.nn.models.trm_module import TRMModule\n\ndef decode_token(t):\n    if t \u003c= 2: return 0\n    return int(t-2)\n\ndevice = 'mps' if torch.backends.mps.is_available() else 'cpu'\nckpt = 'train/runs/\u003ctimestamp\u003e/checkpoints/last.ckpt'  # replace with your checkpoint\nmodel = TRMModule.load_from_checkpoint(ckpt, map_location=device).to(device).eval()\n\nN=3; GRID=4; MAX=6\ndm = SudokuDataModule(batch_size=N, num_workers=0, grid_size=4, max_grid_size=6,\n                      generate_on_fly=True, num_train_puzzles=2000, num_val_puzzles=N)\ndm.setup('fit')\nbatch = next(iter(dm.val_dataloader()))\nfor k,v in list(batch.items()):\n    if isinstance(v, torch.Tensor): batch[k]=v.to(device)\n\nwith torch.no_grad():\n    carry = model.initial_carry(batch)\n    while True:\n        carry, outputs = model.forward(carry, batch)\n        if carry.halted.all(): break\n    preds = outputs['logits'].argmax(dim=-1)\n\nlabels = batch['output']\nmask = labels != -100\npix_acc = (preds==labels)[mask].float().mean().item()\nprint(f'Batch pixel accuracy: {pix_acc:.3f}\\n')\n\nfor i in range(N):\n    pred = preds[i].cpu().numpy().reshape(MAX,MAX)\n    lab  = labels[i].cpu().numpy().reshape(MAX,MAX)\n    inp  = batch['input'][i].cpu().numpy().reshape(MAX,MAX)\n    pred4 = np.vectorize(decode_token)(pred)[:GRID,:GRID]\n    lab4  = np.vectorize(decode_token)(lab)[:GRID,:GRID]\n    inp4  = np.vectorize(decode_token)(inp)[:GRID,:GRID]\n    exact = int((pred4==lab4).all())\n    print(f'Sample {i+1} (exact: {exact})')\n    print('Input:')\n    print(inp4)\n    print('Target:')\n    print(lab4)\n    print('Pred:')\n    print(pred4)\n    print('-'*28)\nPY\n```\n\nExample result from our 60‑epoch stable run (128×2, N_supervision=2): all samples exact and batch pixel accuracy 1.000.\n\nExamples (from a held‑out batch)\n\nInput:\n```\n[[0 0 0 0]\n [0 4 0 1]\n [0 0 0 4]\n [0 3 1 0]]\n```\nTarget / Pred:\n```\n[[1 2 4 3]\n [3 4 2 1]\n [2 1 3 4]\n [4 3 1 2]]\n```\n\nInput:\n```\n[[0 0 4 2]\n [2 0 0 1]\n [3 0 0 4]\n [0 0 0 0]]\n```\nTarget / Pred:\n```\n[[1 3 4 2]\n [2 4 3 1]\n [3 1 2 4]\n [4 2 1 3]]\n```\n\nInput:\n```\n[[0 4 1 3]\n [0 0 0 0]\n [4 0 3 1]\n [0 1 0 0]]\n```\nTarget / Pred:\n```\n[[2 4 1 3]\n [1 3 2 4]\n [4 2 3 1]\n [3 1 4 2]]\n```\n\n# Note to contributors\n\nIf you would like to make contributions to this codebase, here are things you can do:\n\n- Help reproduce the results of the original paper\n- Implement missing features (carry, puzzle embeddings)\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fallthingssecurity%2Ftrm_sudoku","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fallthingssecurity%2Ftrm_sudoku","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fallthingssecurity%2Ftrm_sudoku/lists"}