feat: connect compositional scaling tasks

1 files changed, 46 insertions, 0 deletions

diff --git a/experiments/synthetic_smoke.py b/experiments/synthetic_smoke.py
new file mode 100644
index 0000000..c118813
--- /dev/null
+++ b/experiments/synthetic_smoke.py
@@ -0,0 +1,46 @@
+"""Fast checks for the synthetic depth-scaling task plumbing."""
+import os
+import sys
+
+import torch
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from experiments.run import build, get_args, load_task
+
+
+def check_task(dataset, expected_in, expected_out, extra):
+    argv = [
+        "run.py", "--dataset", dataset, "--device", "cpu",
+        "--task_train_examples", "64", "--task_test_examples", "32",
+        "--batch_size", "16", "--depth", "2", "--width", "12",
+    ] + extra
+    old = sys.argv
+    try:
+        sys.argv = argv
+        args = get_args()
+    finally:
+        sys.argv = old
+    train, test, n_in, n_out = load_task(args, "cpu")
+    assert (n_in, n_out) == (expected_in, expected_out)
+    args.n_in, args.n_out = n_in, n_out
+    x, y = next(iter(train))
+    assert x.shape == (16, n_in)
+    assert y.min() >= 0 and y.max() < n_out
+    for mode in ("bp", "dfa", "sdil"):
+        args.mode = mode
+        net, _ = build(args, "cpu")
+        assert net.logits(x).shape == (16, n_out)
+    assert sum(batch_y.numel() for _, batch_y in test) == 32
+    print(f"{dataset}: input={n_in}, classes={n_out}, fixed task seed={args.task_seed}")
+
+
+if __name__ == "__main__":
+    torch.manual_seed(0)
+    check_task("teacher", 20, 4,
+               ["--task_n_in", "20", "--task_classes", "4",
+                "--teacher_depth", "3", "--teacher_width", "10"])
+    check_task("hierarchical", 16, 4,
+               ["--task_levels", "4", "--task_classes", "4"])
+    check_task("tentmap", 5, 2,
+               ["--task_levels", "4", "--task_n_in", "5"])
+    print("ALL SYNTHETIC TASK CHECKS PASSED")