diff --git a/CLAUDE.md b/CLAUDE.md index 085ef6f..518cbed 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -75,23 +75,23 @@ The repo's `uv.toml` exists solely to make that install work: the user's global **Suite tiers.** The full suite is ~10 min, dominated by the large-scale milestone proofs tagged `# suite-tier: full` (`test_grid_context_selfmod`, `test_grid_nbhd_selfmod`, `test_grid_countmap_selfmod`, `test_delta_selfmod`, `test_composed_generalization`). `./esper fast` (= `run_tests.sh fast`) skips those for a ~2 min local gate; `./esper suite` (default `full`) runs everything and is what CI runs. **The fast tier is a strict subset at FULL budget** — no reduced iterations, no relaxed thresholds — so it never makes a weakened claim; it still exercises every code path (structural, the ES operator fit, and the self-mod `meta_fit_selfmod` core via `test_selfmod_memory`) and only defers the large-scale milestone proofs. A test opts *out* of `fast` by tagging a `# suite-tier: full` comment line at its top (grepped by `run_tests.sh`); untagged tests default to `fast`, so new tests run in both unless they declare themselves heavy. -`run_tests.sh [full|fast]` generates sample fixtures (a `.bin` grid and a `.task` bundle), runs the tier's `tests/test_*.mojo` with `mojo run -I src ` (the `-I src` puts the source modules on the import path), runs the `src/main.mojo` driver, and finally generates a few `.task` bundles via `synth_tasks.generate_task_groups` and runs the `src/arc_solve.mojo` held-out generalization driver over them. To run a single test directly: `mojo run -I src tests/test_demo_fitness.mojo`. Tests `raise Error(...)` on assertion failure rather than using a test framework — follow that pattern for new tests. CI (`.github/workflows/ci.yml`) also enforces `mojo format` (note: this build's `mojo format` has **no `--check` flag** — CI runs the formatter then checks for a git diff). +`run_tests.sh [full|fast]` generates sample fixtures (a `.bin` grid and a `.task` bundle), runs the tier's `tests/test_*.mojo` with `mojo run -I src ` (the `-I src` puts the source modules on the import path), runs the `src/main.mojo` driver, and finally generates a few `.task` bundles via `synth_tasks.generate_task_groups` (plus, full tier only, one shape-changing bundle via `generate_shape_task_groups` so the driver's shape dispatch runs in CI) and runs the `src/arc_solve.mojo` held-out generalization driver over them. To run a single test directly: `mojo run -I src tests/test_demo_fitness.mojo`. Tests `raise Error(...)` on assertion failure rather than using a test framework — follow that pattern for new tests. CI (`.github/workflows/ci.yml`) also enforces `mojo format` (note: this build's `mojo format` has **no `--check` flag** — CI runs the formatter then checks for a git diff). ## Architecture map - `src/hope.mojo` — Core data structures + operator *execution* (no learning, so it has no `arc_io`/ES deps): `ArcGrid` (owned row-major Float32 grid), `HopeArena` (move-only bump allocator with `bump[T](count)` / `alloc_node[T]`), the **POD** `HopeNode` (raw `slow`/`fast` weight slices into an arena + inline child indices) and its `build_node` factory, and the **structured learned operator** — `OP_DIM` layout (6-param centered affine + 10-entry **normalized** color LUT, stored /9 to keep colour at the affine's ~unit scale), `seed_identity_operator`, and `apply_operator`. The operator is a **bilinear geometry gather** where each of the four corner input cells is mapped through the colour LUT (`_color_of`) *before* the blend — **colour-then-gather** decouples the colour fit from the geometry's precision. Smooth so the ES has a real gradient everywhere, yet exact at integer params. The generic demo containers `ExamplePair[E]` and `Task[E]` live here, with `ArcTaskPair`/`ArcTask` as grid `comptime` aliases (the ES core is generic over the Example type, so the containers must carry it — Mojo checks generic bodies eagerly). The old `prim_*`/`apply_primitive` and `update_fast_weights` remain but are **dormant** (off the operator path); `forward_with_learning` moved to `esper_evolution.mojo` (import-DAG reasons). -- `src/esper_evolution.mojo` — All learning, **generic over `[M: Memory]`** (Phase B): `fitness[M]`, `evolve_fast_weights[M]`, `fit_operator[M]`, `reptile_meta_train[M]`, and `ESWorkspace[M]` are parametric; their SIMD/FMA bodies are unchanged. `fitness[M]` runs a candidate memory (`M.apply`) on each demo input, scores the Domain metric (`M.Dom.distance`) vs the demo output (penalizing a heavy constant if output area ≠ input area — same-shape), minus an L2 anchor toward the slow prior. It replaced the deleted `evaluate_primitives` memorization surrogate. `ESWorkspace` holds **param-sized** ES vectors plus a **grid-sized** `op_output` scratch and a per-parameter `scale` (the colour group gets `COLOR_SCALE` < 1 — a diagonal **preconditioner** so one step size fits both geometry and colour). `evolve_fast_weights` is one antithetic-sampling ES step over the demos — **real Gaussian noise** (`randn_float64`), mirrored `grad += (F(w+σ·scale·ε) − F(w−σ·scale·ε))·ε`, then `W_fast += alpha/(2Nσ)·scale·grad`. `fit_operator` is the **annealed** fit loop (the shared `FIT_*` schedule: sigma 0.5→0.01, alpha 0.1→0.003 over ~4000 iters — explore wide, then settle onto exact integers; the wide `FIT_SIGMA0` is what makes transpose's four-param move robust). `forward_with_learning` ties it together: fit `node.fast` to the demos (anchored to `node.slow`), then `apply_operator` on the test input. **M9 (the second timescale)**: `reptile_meta_train` is the outer, low-frequency loop that turns `slow` from a fixed identity anchor into a *meta-learned prior* — for each task it fits `fast` in-context FROM the current prior (reusing `fit_operator`) then Reptile-nudges `slow += META_LR·(fast − slow)` (`reptile_update`/`copy_weights` are the SIMD/FMA weight helpers). Crucially the **eval** schedule is a *narrow exploit* fit (`EVAL_SIGMA0=0.12`, `EVAL_ITERS=300`), not the wide `FIT_SIGMA0` — a warm start only helps when you can then afford a cheap local fit (the wide sigma explores the whole space and washes out the init). Meta-train keeps the wide sigma (it must *discover* the operator from cold). Also hosts `meta_fit_selfmod[M: SelfModMemory]` (the self-mod family's meta-fit) and `fit_geomcolor` (block 5: colour write → pre-map the demo list through V once per task → the standard `fit_operator[AttnGatherMemory]` on the composed state's 7 attention slots — zero core change, hot loop stays alloc-free); `fit_geomcount` is the same shape for the count-content module; both drivers apply **constant-compute budgeting** — iters × `FIT_DEMO_REF`/n_demos, so every task gets the same total demo-evaluations and the n=8 proofs are unchanged). +- `src/esper_evolution.mojo` — All learning, **generic over `[M: Memory]`** (Phase B): `fitness[M]`, `evolve_fast_weights[M]`, `fit_operator[M]`, `reptile_meta_train[M]`, and `ESWorkspace[M]` are parametric; their SIMD/FMA bodies are unchanged. `fitness[M]` runs a candidate memory (`M.apply`) on each demo input, scores the Domain metric (`M.Dom.distance`) vs the demo output (penalizing a heavy constant if output area ≠ input area — same-shape), minus an L2 anchor toward the slow prior. It replaced the deleted `evaluate_primitives` memorization surrogate. `ESWorkspace` holds **param-sized** ES vectors plus a **grid-sized** `op_output` scratch and a per-parameter `scale` (the colour group gets `COLOR_SCALE` < 1 — a diagonal **preconditioner** so one step size fits both geometry and colour). `evolve_fast_weights` is one antithetic-sampling ES step over the demos — **real Gaussian noise** (`randn_float64`), mirrored `grad += (F(w+σ·scale·ε) − F(w−σ·scale·ε))·ε`, then `W_fast += alpha/(2Nσ)·scale·grad`. `fit_operator` is the **annealed** fit loop (the shared `FIT_*` schedule: sigma 0.5→0.01, alpha 0.1→0.003 over ~4000 iters — explore wide, then settle onto exact integers; the wide `FIT_SIGMA0` is what makes transpose's four-param move robust). `forward_with_learning` ties it together: fit `node.fast` to the demos (anchored to `node.slow`), then `apply_operator` on the test input. **M9 (the second timescale)**: `reptile_meta_train` is the outer, low-frequency loop that turns `slow` from a fixed identity anchor into a *meta-learned prior* — for each task it fits `fast` in-context FROM the current prior (reusing `fit_operator`) then Reptile-nudges `slow += META_LR·(fast − slow)` (`reptile_update`/`copy_weights` are the SIMD/FMA weight helpers). Crucially the **eval** schedule is a *narrow exploit* fit (`EVAL_SIGMA0=0.12`, `EVAL_ITERS=300`), not the wide `FIT_SIGMA0` — a warm start only helps when you can then afford a cheap local fit (the wide sigma explores the whole space and washes out the init). Meta-train keeps the wide sigma (it must *discover* the operator from cold). Also hosts `meta_fit_selfmod[M: SelfModMemory]` (the self-mod family's meta-fit) and `fit_geomcolor` (block 5: colour write → pre-map the demo list through V once per task → the standard `fit_operator[AttnGatherMemory]` on the composed state's 7 attention slots — zero core change, hot loop stays alloc-free); `fit_geomcount` is the same shape for the count-content module; both drivers apply **constant-compute budgeting** — iters × `FIT_DEMO_REF`/n_demos, so every task gets the same total demo-evaluations and the n=8 proofs are unchanged). **Shape seam (Next #1)**: `fitness_shape[M: ShapeMemory]`/`fit_shape[M]` are the shape-aware generic siblings (the memory predicts its own output dims; a predicted/true area mismatch is heavy-penalized), and `fit_shape_geom` is the per-task driver — a **two-start** fit (the same two frame seeds every task, from the written slope; winner by demo fitness — an honest multi-start, no staging), each start DISCOVER (soft, temperature searched) then SETTLE (temperature hard-frozen at `SHAPE_BETA_READ` via the settle-variant type, sigma held at `SHAPE_SIGMA_FLOOR` — the sharp landscape's staircase step scale, below which ES updates are pure noise — alpha decayed), constant-compute budgeted like the rest. - `src/arc_io.mojo` — `_read_grid_block` reads one validated grid-block (16-byte header: two little-endian Int64 `rows`/`cols`, then the float32 payload) at an offset, advancing it; shared by `load_arc_grid` (single `.bin`) and `load_arc_task` (a `.task` **bundle**: `[n_train][n_test]` then the train/test grid-blocks → an `ArcTask`). Header lengths are validated against the file so truncated/malformed input raises. `calculate_fitness` = SIMD negative MSE (continuous ES signal); `exact_match` = discrete reward (fraction equal after `round`). (Note: Mojo's `open` only accepts `r/w/rw/a`; `read_bytes()` returns raw bytes.) **Phase B:** the `Domain` trait + `GridDomain` live here — a Domain is `(associated Example type, distance, score, capacity)`; `GridDomain.Example = ArcGrid` and its metrics wrap `calculate_fitness`/`exact_match`. The ES core reaches metrics only through a Domain, never ARC directly. -- `src/memory.mojo` — **The memory seam** (traits only; the families live in the `memory_*` modules below, all flat `-I src` modules with direct imports — no re-export façade). `Memory`: what the ES fits in-context — `param_dim` (replaces `OP_DIM`), `seed`, `fill_scale` (the per-memory ES preconditioner), `apply(weights, inp: Self.Dom.Example, dst)`, with an associated `comptime Dom: Domain` reached as `M.Dom` (Mojo traits can't take parameters, so the domain is an associated type, not `Memory[D]`). `SelfModMemory`: the self-write counterpart — `slow_dim`/`state_dim`/`seed_slow`/`fill_scale`/`adapt` (writes the fast state from the demos via the memory's own rule)/`apply`. No runtime memory-selector — the memory is a compile-time choice, each measured on the subset it expresses. -- `src/memory_es.mojo` — The **ES-fit forward family**: `OperatorMemory` (the structured Phase-A affine+LUT operator, **dormant** — subsumed emergently by `GeomColorComposedMemory`, kept only as the arc_solve/M8 baseline; owns `COLOR_SCALE`), `MLPMemory` (**B1**, the first emergent memory: per-cell `1→H→1` tanh MLP, learns recolor with no hand-coded LUT, output squashed to `[0,9]`), `SeqOperatorMemory`/`SeqMLPMemory` (**B2**, the sequence-domain pair proving the seam is domain-generic), `AttnGatherMemory` (**B3** emergent geometry: a 7-param learned position-attention gather — `M`, `t`, temperature `beta=raw²`; integer projections ⇒ exact flip/transpose; the softmax is **windowed** — a `(2·ATTN_WINDOW+1)²` window centred on `q`, bit-identical on synth-scale grids and ~3-5× cheaper at real-ARC 30×30). -- `src/memory_composed.mojo` — `GeomColorComposedMemory` (**block 5** — the emergent retirement of `OperatorMemory`): a count-signature colour table **written closed-form** from the demos (geometry-invariant — counts are position-free under permutation geometry) composed with the AttnGather ES run on V-pre-mapped demos (**colour-then-gather**; the fit driver is `fit_geomcolor` in esper_evolution). `fill_scale` zeroes the V group so no ES path can move the written table. The writes are few-demo hardened: global-min greedy INJECTIVE assignment with identity defaults for unseen colours and identity preference on exact ties (the corpus median is 3 demos). Also `GeomCountComposedMemory` (the content×geometry block): the same recipe one level up — a neighbourhood-count content rule `out = geom(M(count_P(in)))`, its (P, M) written **geometry-invariantly from histogram signatures** (count rules commute with the lattice-symmetry geometry class), geometry via `fit_geomcount` on content-premapped demos. +- `src/memory.mojo` — **The memory seam** (traits only; the families live in the `memory_*` modules below, all flat `-I src` modules with direct imports — no re-export façade). `Memory`: what the ES fits in-context — `param_dim` (replaces `OP_DIM`), `seed`, `fill_scale` (the per-memory ES preconditioner), `apply(weights, inp: Self.Dom.Example, dst)`, with an associated `comptime Dom: Domain` reached as `M.Dom` (Mojo traits can't take parameters, so the domain is an associated type, not `Memory[D]`). `SelfModMemory`: the self-write counterpart — `slow_dim`/`state_dim`/`seed_slow`/`fill_scale`/`adapt` (writes the fast state from the demos via the memory's own rule)/`apply`. `ShapeMemory`: the output-size seam (Next #1) — `write` (infers the shape rule closed-form from the demo dim-pairs), `out_rows`/`out_cols` (predict the output dims), `apply(state, inp, out_rows, out_cols, dst)`; a distinct trait so the same-shape core and memories stay untouched. No runtime memory-selector — the memory is a compile-time choice, each measured on the subset it expresses. +- `src/memory_es.mojo` — The **ES-fit forward family**: `OperatorMemory` (the structured Phase-A affine+LUT operator, **dormant** — subsumed emergently by `GeomColorComposedMemory`, kept only as the arc_solve/M8 baseline; owns `COLOR_SCALE`), `MLPMemory` (**B1**, the first emergent memory: per-cell `1→H→1` tanh MLP, learns recolor with no hand-coded LUT, output squashed to `[0,9]`), `SeqOperatorMemory`/`SeqMLPMemory` (**B2**, the sequence-domain pair proving the seam is domain-generic), `AttnGatherMemory` (**B3** emergent geometry: a 7-param learned position-attention gather — `M`, `t`, temperature `beta=raw²`; integer projections ⇒ exact flip/transpose; the softmax is **windowed** — a `(2·ATTN_WINDOW+1)²` window centred on `q`, bit-identical on synth-scale grids and ~3-5× cheaper at real-ARC 30×30; `apply_shaped` decouples the query grid (output dims) from the source grid). Also `attn_gather_toroidal` (the shape-seam gather): the output-shaped read with a **toroidal source** (wrapped images — makes tiling's sawtooth expressible), an **extent-relative translation** `trel` (absorbs tiling's size-dependent phase), and the query **normalized by the written shape slope** (resize-as-identity — content never re-learns the scale the shape rule knows). `AttnGatherMemory` itself is untouched by the shape work (same-shape proofs bit-identical). +- `src/memory_composed.mojo` — `GeomColorComposedMemory` (**block 5** — the emergent retirement of `OperatorMemory`): a count-signature colour table **written closed-form** from the demos (geometry-invariant — counts are position-free under permutation geometry) composed with the AttnGather ES run on V-pre-mapped demos (**colour-then-gather**; the fit driver is `fit_geomcolor` in esper_evolution). `fill_scale` zeroes the V group so no ES path can move the written table. The writes are few-demo hardened: global-min greedy INJECTIVE assignment with identity defaults for unseen colours and identity preference on exact ties (the corpus median is 3 demos). Also `GeomCountComposedMemory` (the content×geometry block): the same recipe one level up — a neighbourhood-count content rule `out = geom(M(count_P(in)))`, its (P, M) written **geometry-invariantly from histogram signatures** (count rules commute with the lattice-symmetry geometry class), geometry via `fit_geomcount` on content-premapped demos. And the **shape seam pair** (Next #1): `ShapeGeomComposedMemory` — layout `[0:7]` attention | `[7:9]` trel | `[9:13]` shape rule (`out = round(k·in + b)` per axis, least-squares-written, frozen) — whose gather is `attn_gather_toroidal`; plus `ShapeGeomSettleMemory`, a thin delegating variant whose only difference is a frozen temperature slot (`fill_scale` is static per type, so the fit driver's settle phase is a type). A k-fold size change has **two canonical identity frames** — rescaled (`M=I`) and periodic (`M=kI, trel=(k−1)/2`, via `seed_periodic`) — which is why the fit driver is a two-start (see `fit_shape_geom`). - `src/memory_selfmod.mojo` — The **self-write family, core mechanisms** (B4): `RecolorSelfWrite` (fixed-projection checkpoint), `RecolorSelfModMemory` (meta-learned associative read; a fresh recolor adapts in ONE pass), `DeltaSelfModMemory` (gated delta-rule self-write `S ← (1−α)S + η(v−S·k)k` with self-generated key/η/α, sequence domain). Fast adaptation is the memory's own write rule over the demos; the ES meta-learns only the small slow vector. - `src/memory_selfmod_grid.mojo` — The **2-D grid self-mod memories** (ARC-AGI-2 blocks 1–4): `GridContextSelfModMemory` (additive centre/neighbour rules via outer-product keys), `GridNbhdSelfModMemory` (the disjunctive/count class: centre-free Moore-8 histogram key + sigmoid-threshold read, whole 2-level rule inferred in-context), `GridCountMapSelfModMemory` (arbitrary count→colour maps: meta-learned scoring salience + soft count-bin value table). - `src/main.mojo` — End-to-end driver: builds an `OP_DIM` node, seeds slow (prior) and fast (init) to identity, learns `flip_h` in-context via `forward_with_learning`, prints result + held-out exact match. `mojo run -I src src/main.mojo`. -- `src/arc_solve.mojo` — **Held-out generalization driver** (replaced `benchmark.mojo`). Takes `.task` bundle paths via argv (shell-globbed), fits each task's **emergent composed memory** (`GeomColorComposedMemory` via `fit_geomcolor` — the operator's successor) on its train pairs, scores the **unseen** test pair(s), and reports per-task held-out + train-fit + the train/test gap, then the aggregate solve rate. A pair whose output area ≠ input area honestly scores 0 (the operator is same-shape — this guards against an OOB compare on real shape-changing ARC tasks). Raises on 0 solved (a CI regression signal for the synth bundles) **unless** the first arg is `--report` (honest real-ARC eval mode, where 0% is a legitimate number). `--fit N ITERS` overrides the ES budget (default = the full proven `FIT_*`): the real-corpus runs use a smaller **documented** budget, quoted with the number, because full-budget fits at 30×30 ARC scale are compute-prohibitive; when EVERY test pair is shape-changing the fit is skipped (held-out is 0 by construction — exact, saves ~1/3 of corpus compute). Seeds the RNG **per task** (`SOLVE_SEED`, inside `solve_task`) so each task's stochastic ES fit depends only on the task, not its position in argv — the benchmark number is reproducible and invariant to task ordering / sharding. Uncheatable by memorization. **`eval_parallel.sh`** (repo root) shards a `.task` directory round-robin across `nproc` worker processes (optional trailing `fit_N fit_iters` pair forwards the corpus budget) (one driver invocation per shard) and re-aggregates the per-task lines — ~`nproc`× faster, identical numbers (process-level sharding, since the `ESWorkspace`/global RNG aren't thread-safe for in-process parallelism). +- `src/arc_solve.mojo` — **Held-out generalization driver** (replaced `benchmark.mojo`). Takes `.task` bundle paths via argv (shell-globbed), fits each task's **emergent composed memory** on its train pairs, scores the **unseen** test pair(s), and reports per-task held-out + train-fit + the train/test gap + a trailing `mem: same|shape` marker (appended AFTER the existing fields — `eval_parallel.sh` reads held-out positionally), then the aggregate solve rate. **Dispatch** (rung b): any train pair whose DIMS differ routes to `ShapeGeomComposedMemory`/`fit_shape_geom` (the memory predicts its output dims from the written rule; a predicted/true dims mismatch scores that pair 0, never applied — no OOB); all-same-dims keeps the byte-identical `GeomColorComposedMemory`/`fit_geomcolor` path, where a test pair whose output area ≠ input area honestly scores 0. Driver-level routing on a closed-form observable of the demos, not a memory-selector — the same-shape memory provably scores 0 on the dispatched class. Raises on 0 solved (a CI regression signal for the synth bundles) **unless** the first arg is `--report` (honest real-ARC eval mode, where 0% is a legitimate number). `--fit N ITERS` overrides the ES budget (default = the full proven `FIT_*`): the real-corpus runs use a smaller **documented** budget, quoted with the number, because full-budget fits at 30×30 ARC scale are compute-prohibitive; for same-shape-dispatched tasks where EVERY test pair is shape-changing the fit is skipped (held-out is 0 by construction — exact). Seeds the RNG **per task** (`SOLVE_SEED`, inside `solve_task`) so each task's stochastic ES fit depends only on the task, not its position in argv — the benchmark number is reproducible and invariant to task ordering / sharding. Uncheatable by memorization. **`eval_parallel.sh`** (repo root) shards a `.task` directory round-robin across `nproc` worker processes (optional trailing `fit_N fit_iters` pair forwards the corpus budget) (one driver invocation per shard) and re-aggregates the per-task lines — ~`nproc`× faster, identical numbers (process-level sharding, since the `ESWorkspace`/global RNG aren't thread-safe for in-process parallelism). - `tools/arc_compiler.py` — The one sanctioned Python component: offline compiler. `_write_grid`/`_save_grid` (single grid `.bin`) and `_save_task` (a `.task` bundle) are the single source of the on-disk formats; `compile_arc_json` converts ARC JSON to per-grid `.bin`s. `compile_task_to_bundle`/`compile_arc_dir` + a `__main__` CLI (`python tools/arc_compiler.py `) batch-ingest a real **ARC-AGI 2** corpus directory into `{task_id}.task` bundles for `arc_solve.mojo` (the M8 path). The corpus is **not** vendored (`data_bin/` and the `arg-agi-2-data` symlink are gitignored). Not part of the runtime path. - `tools/synth_tasks.py` — Offline deterministic generator: `generate_tasks` (single-grid pairs) and `generate_task_groups` (ARC-shaped `.task` bundles: N train demos + a held-out test, per transform), reusing the `arc_compiler` writers. It is the *ground-truth generator* the engine must rediscover — the symbolic transforms (flip/transpose/recolor/shift) live here, never in the engine. -- `tests/` — `test_arena`, `test_operator` (hand-set weights reproduce the transforms exactly), `test_fitness`, `test_demo_fitness` (keystone: ES fits `flip_h` and generalizes), `test_forward_learning` (end-to-end node path, fit-once/generalize-many), `test_task_loader` (bundle round-trip), `test_shape` (same-shape fits; shape-change penalized, no crash), `test_generalization` (**whole expressible subset** learned to ≥0.95 held-out), `test_meta_prior` (**M9**: a Reptile-meta-learned `slow` prior fits a fresh flip_h task to ≥0.95 held-out at a narrow eval budget where a cold identity prior scores ~0), `test_mlp_memory` (**B1**: the emergent `MLPMemory` learns recolor to ≥0.95 held-out through the generic seam, no LUT), `test_composed_generalization` (**block 5, the OperatorMemory retirement proof**: the composed memory matches the operator's whole subset ≥0.95 held-out cold AND solves a composed flip∘recolor no single memory expresses; geometry-only ablation control must fail), `test_composed_content` (**content×geometry**: `geom∘countmap` solved 1.0 cold; the correspondence-statistic and content-ablation controls must fail), `test_few_demo` (**few-demo robustness** at the corpus-median 3 demos: aggregate ≥0.85 + an n=8 regression guard), `test_io`, plus the B2–B4 family tests (`test_seq_domain`, `test_attn_memory`, `test_selfmod_memory`, `test_delta_selfmod`, `test_grid_context_selfmod`, `test_grid_nbhd_selfmod`, `test_grid_countmap_selfmod`). All import the real `src` modules via `-I src`. Phase-A expressible subset = {identity, flip_h, flip_v, transpose, recolor}; `shift` deferred (the affine zero-fills, synth `_shift` wraps). The ES-based tests anneal a few thousand iters; the full suite is ~10 min (dominated by the `# suite-tier: full`-tagged milestone proofs), or ~2 min via `./esper fast` — see "Testing". +- `tests/` — `test_arena`, `test_operator` (hand-set weights reproduce the transforms exactly), `test_fitness`, `test_demo_fitness` (keystone: ES fits `flip_h` and generalizes), `test_forward_learning` (end-to-end node path, fit-once/generalize-many), `test_task_loader` (bundle round-trip), `test_shape` (same-shape fits; shape-change penalized, no crash), `test_generalization` (**whole expressible subset** learned to ≥0.95 held-out), `test_meta_prior` (**M9**: a Reptile-meta-learned `slow` prior fits a fresh flip_h task to ≥0.95 held-out at a narrow eval budget where a cold identity prior scores ~0), `test_mlp_memory` (**B1**: the emergent `MLPMemory` learns recolor to ≥0.95 held-out through the generic seam, no LUT), `test_composed_generalization` (**block 5, the OperatorMemory retirement proof**: the composed memory matches the operator's whole subset ≥0.95 held-out cold AND solves a composed flip∘recolor no single memory expresses; geometry-only ablation control must fail), `test_composed_content` (**content×geometry**: `geom∘countmap` solved 1.0 cold; the correspondence-statistic and content-ablation controls must fail), `test_few_demo` (**few-demo robustness** at the corpus-median 3 demos: aggregate ≥0.85 + an n=8 regression guard), `test_shape_change` (**the output-size seam**: {crop1, flip_h_crop1, subsample2, upscale2, tile2} each ≥0.95 held-out at a FRESH input size, per-task cold — demos drawn at varying sizes so the shape rule is identifiable; controls: no-shape-write → 0, plain non-toroidal gather fails tile2), `test_io`, plus the B2–B4 family tests (`test_seq_domain`, `test_attn_memory`, `test_selfmod_memory`, `test_delta_selfmod`, `test_grid_context_selfmod`, `test_grid_nbhd_selfmod`, `test_grid_countmap_selfmod`). All import the real `src` modules via `-I src`. Phase-A expressible subset = {identity, flip_h, flip_v, transpose, recolor}; `shift` deferred (the affine zero-fills, synth `_shift` wraps). The ES-based tests anneal a few thousand iters; the full suite is ~10 min (dominated by the `# suite-tier: full`-tagged milestone proofs), or ~2 min via `./esper fast` — see "Testing". ## Conventions to preserve when extending diff --git a/docs/JOURNAL.md b/docs/JOURNAL.md index d3f5b8e..bcc1bc7 100644 --- a/docs/JOURNAL.md +++ b/docs/JOURNAL.md @@ -956,3 +956,323 @@ fought). **d511f180 — the block's real-corpus exhibit — recovered: held-out v2 corpus budget (was 0.75/0.80; M8's full-budget ES-fit LUT had solved it, the un-hardened write had lost it). `test_few_demo` (full tier, ~63s) locks the bars: n=3 aggregate ≥0.85 (measured 0.87; pre-hardening 0.74), n=8 per-family ≥0.95. + +--- + +## 2026-07-03 13:12 — Shape change: the output-size seam + first shape-change family (Vision A, Next #1) + +The roadmap's Next #1, and the binding constraint the 07-03 re-measure named: 32% of both corpus +splits change shape (out dims ≠ in dims), and the whole engine was hard-wired same-shape — every +memory's `apply` wrote exactly `inp.rows × inp.cols` cells, `fitness[M]` slammed `-1e9` on any demo +with `in_n ≠ out_n`, and `arc_solve` scored shape-changing pairs 0. Output shape was never even +*represented* as a quantity distinct from the input. This block builds the reusable seam and proves +one shape-change family cold. + +**Why a new trait, not an extension of `Memory`.** A same-shape `Memory.apply(weights, inp, dst)` +has no place to learn a different output size, and giving all 10+ existing memories an output-shape +method (defaulting to input shape) would churn every one and reintroduce the OOB the `-1e9` guard +exists to prevent. So `ShapeMemory` is a distinct trait (parallel to `SelfModMemory`, the B4 +precedent): `write(state, demos)` infers the shape rule closed-form, `out_rows`/`out_cols` predict +the output dims, and `apply(state, inp, out_rows, out_cols, dst)` produces that many cells. The +same-shape core and every existing memory are untouched; the shape-aware fitness/fit driver is a +generic `[M: ShapeMemory]` sibling. Additive, zero regressions. + +**Why the composition pattern again.** A shape-change task factors — like block 5 and the +content×geometry block — into two factors fit on signals invariant to each other: +1. a **shape rule** `out = round(k·in + b)` per axis, WRITTEN closed-form by least-squares over the + demo dim-pairs (position-free shape arithmetic, no geometry knowledge, one pass, never + ES-searched); and +2. the **content** — the *proven* AttnGather gather, generalized so its query grid is the OUTPUT grid + and it reads the INPUT grid. `M=I` reads a centred crop, `M=sI` a subsample, `M=±perm` a + flip/transpose within the resize. ES-fit over the same 7 attention params, on the same B3 + landscape, only reading a differently-sized output. + +**Why an output-shaped gather is a one-variable change.** `AttnGatherMemory.apply` already loops the +query over the grid centred on `inp` and gathers from `inp`. Decoupling the *query* extent/centre +(now the output's) from the *source* extent/centre (still the input's) — plus the output stride on +`dst` — is the whole mechanism. For out==in it reduces to the old code bit-identically, so the new +`apply_shaped` is what `apply` now delegates to; `test_attn_memory`/the fast gate stay 1.0 unchanged. + +**Why the demos must vary input size.** With one input size the shape rule's slope/intercept are +underdetermined (only their combination at that size is pinned) — the honest analogue of the n=2 +signature ties. So the proof draws each demo (and the held-out test) at a RANDOM size in [4,8]: +`≥2` distinct sizes identify `(k,b)`, and the fresh-size test is an uncheatable probe that the *rule* +generalizes, not a memorized size. The least-squares fallback (mean ratio, b=0) keeps the write exact +at a fixed size for the underdetermined case; documented, not fought. + +**Where the ES search lands.** `fill_scale` zeros the shape-rule slots (the GeomColor freeze trick), +so the ES moves only the 7 attention params; the written shape rule rides along frozen. `fitness_shape` +scores at the OUTPUT area and heavy-penalizes a predicted-shape/true-shape area mismatch (the honest +successor to the same-shape guard — a wrong shape rule can't be rescued by content). The L2 anchor's +frozen-slot terms cancel in the antithetic F+−F−, so only content feels it. `fit_shape` is a +self-contained annealed ES (its own scratch, sized to output capacity, like `meta_fit_selfmod`); +`fit_shape_geom` is the per-task driver (write → fit, constant-compute budgeted). + +**Result — `test_shape_change` (full tier), cold, held-out at a fresh size:** +- Ckpt A: the least-squares write recovers crop1's `(k=1, b=−2)` per axis exactly. +- Ckpt B: `{crop1, flip_h_crop1, subsample2}` each **held-out 1.0**, per-task cold — including + subsample2, where the ES had to find `M=2I, t=(−0.5,−0.5)` (a 2× scale-up) from the identity seed. +- Control: the SAME content fit WITHOUT the shape write (identity shape rule) predicts the wrong + output size on every pair → **held-out 0.0** — the inferred shape rule is load-bearing, not + scaffolding. +- Fast gate + `mojo format` clean; same-shape numbers unchanged (bit-identical gather). + +Synth side: `SHAPE_TRANSFORMS` (crop1, flip_h_crop1, subsample2) + `generate_shape_task_groups` +(varies input size across a task's demos); `corpus_stats` reports the emitted bundles as 0% same-shape +(the `.task` format already self-describes per-grid dims — the seam was always consumer-side). + +**Deferred (documented):** upscale/tiling (need a floor/modular gather — a follow-on family on this +same seam; the affine gather provably can't express blocky replication); wiring `arc_solve --report` +to score the real 32% (follow-on — keeps this a clean synth proof); colour composition on top of shape +(a `write_color` pre-map, which commutes cellwise). Latent same-shape assumptions in +`_selfmod_meta_fitness` and the selfmod-grid `adapt` output reads are left as-is (those stay same-shape +families). Base for the next entry: 0e138cf. + +## 2026-07-03 17:12 — Re-ran the real ARC-AGI-2 public-eval split (120 tasks), 6 workers, budget 64/1500 + +Re-measured `arc_solve --report` against `data_bin/arc2_eval` (`eval_parallel.sh data_bin/arc2_eval +scratch/arc2_eval_results.txt 6 64 1500`) at the documented corpus budget, same as every prior +real-corpus number. `arc_solve.mojo` still fits `GeomColorComposedMemory` only (block 5) — the +shape-change seam (`ShapeGeomComposedMemory`/`fit_shape_geom`, landed this session) is **not yet +wired into `arc_solve`** (a documented deferral), so this run is a direct like-for-like comparison to +the prior eval-split baseline, not a measurement of the shape work. + +**Result:** `Solved 0 / 120 (solve rate: 0.000%, mean held-out: 0.387962)`, scored in 8204s (~137 min +wall on 6 workers — noticeably slower than the ~56 min a prior run logged at the same budget; the +machine had only ~1-3GB free RAM for most of the run, so this looks like memory-pressure/swap +slowdown, not a change in per-task cost). + +Mean held-out **0.388** vs the last-logged eval-split number of **0.319** (both 0/120 exact-solved — +the corpus's genuine shape-changing/multi-rule tasks are still out of reach without the shape seam +wired in). The improvement is consistent with noise across the run-to-run seeded ES stochasticity +already documented for these composed memories) rather than a code change on the eval path — no +`arc_solve`/`memory_composed` edits landed between the two measurements. Exact-solve stays 0 because +the composed memory is still same-shape-only on this driver; the shape-change seam is the natural next +rung to wire in before the number can move on the ~32% shape-changing slice this run still can't touch. + +Raw dump: `scratch/arc2_eval_results.txt` (gitignored, reproducible per-task via the per-task RNG +seed). + +## 2026-07-04 07:40 — Upscale/tiling: the output-GROWING families land (Next #1 rung a) — via a measured grid of 12 configurations + +The roadmap's shape rung (a): outputs LARGER than the input — blocky upscale (each cell → an s×s +block) and tiling (the grid replicated k×k). The shape rules (k=s, b=0) were already covered by the +least-squares shape write; the whole battle was the CONTENT gather. This block took ~12 measured +fit-configurations to land honestly; the failures fixed the design, so they're recorded. + +**Correction to the last entry's claim.** "The affine gather provably can't express blocky +replication" is WRONG for upscale: `floor(r/s) = round((r−(s−1)/2)/s)` is an exact identity with no +ties, so nearest-cell reading of an affine map expresses blocky upscale exactly (probe-verified: +hand-set `M = I/s` scores 1.0 at sharp temperature, both parities, s∈{2,3}). What is genuinely +outside any affine `(M, t)` is TILING — `out[r] = in[r mod n]` is a sawtooth. And even for upscale +the EXPRESSIBILITY was never the issue — the FIT was. + +**The mechanism (memory side).** `attn_gather_toroidal` (memory_es) — the output-shaped AttnGather +read with three additions, used by `ShapeGeomComposedMemory.apply`: +1. **Toroidal source**: per-axis displacements wrap into (−extent/2, extent/2], so a query past the + edge reads the input's periodic image — tiling's sawtooth is the nearest WRAPPED cell of an + affine map. A substrate choice (precedent: the selfmod-grid memories' toroidal neighbourhoods), + not a task primitive. Window span capped at the torus period (else a cell is scanned twice). +2. **Extent-relative translation** `trel` (2 new slots, SHAPEGEOM_DIM 11→13): the centred frames + leave tiling with a size-dependent phase (n/2 for k=2) that no constant t can cancel across + varying demo sizes; `q += trel·extent` absorbs it with one size-free parameter. +3. **Query normalization by the WRITTEN shape slope** (`v_out/k`): resize-as-identity. Without it a + resize family's `M = 1/k` and its exactness tolerance shrinks with the output extent (upscale-2's + m11 needed ±0.045 — under the ES's settling noise at any workable sigma; measured: fits parked at + 0.554, the plateau EDGE, four runs in a row). Normalized, `M = I` IS every pure resize and all + tolerances are size-free. The shape factor informing the content factor's coordinate frame is the + composition pattern once more. + +**The fit (driver side) — what the 12 configs taught.** The permutation families read at integer +positions (temperature-insensitive); upscale reads at ±1/4 offsets and NEEDS a sharp read. The +failure grid, each cell measured (fit_shape_geom): +- Temperature SEARCHED (any seed, any preconditioner): the ES drives it SOFT — it optimizes the + Gaussian-smoothed objective, where a soft read is robust to the sampler's own jitter. Geometry + converges exactly; the read blurs; held-out ~0.12. Seeded sharp (raw 3.0) it re-softens to 1.4. +- Temperature ANNEALED up mid-fit (soft→sharp alongside sigma/alpha, wide or narrow): the staircase + landscape + shrinking sigma = a noise walk; the geometry itself diverges (t drifted to ±1.5). +- Temperature FROZEN sharp, sigma annealed to 0.01/0.05: below the staircase's ~0.25 step scale the + antithetic differences are almost always zero — divergence (M diag hit 2.24). +- Temperature FROZEN sharp, sigma floored at 0.15 (the step scale): tile2 1.0/1.0 — the DISCRETE + regime works when the ES is run as the stochastic hill-climber it then is. But upscale parked at + the plateau edge (its pre-normalization tolerance was under the floor), and one seed frame cannot + serve both families (below). +- An L2-anchor-parks-at-plateau-edge hypothesis was falsified cleanly (reg=0 reproduced the same + trajectory to 5 decimals — the 1e-4 anchor is ~2e-6 in fitness, negligible). +**Two identity frames.** With the normalized query, upscale's solution is the SEED (`M = I`) — and +tiling's moved to `M = kI, trel = (k−1)/2` (corner-aligned periodic read), one unit of travel away; +fits from the wrong frame reliably fall into a degenerate constant-read basin (measured both +directions: whichever family's solution is at the seed solves, the other collapses). A k-fold size +change simply HAS two canonical identity continuations — the rescaled plane and the periodic plane — +both derivable from the WRITTEN slope. So `fit_shape_geom` runs the SAME TWO cold starts for every +task, each DISCOVER (wide soft anneal, temperature searched — the proven-smooth landscape) then +SETTLE (temperature hard-frozen at `SHAPE_BETA_READ` via `ShapeGeomSettleMemory` — fill_scale is +static per type, so the phase difference is a thin delegating type — sigma HELD at the 0.15 step +floor where the plateau-edge gradient stays alive, alpha decayed: measured to centre t/trel to +~0.01), inside the same total budget (half each), winner by demo fitness at the hard read. +An honest multi-start — selection by the task's own train signal; no task-specific staging anywhere. +`trel` gets a small ES scale (0.2): it multiplies the extent, and the frame seeds already place it +at its solution — refinement only. + +**Result — `test_shape_change` (full tier, 4m48s):** all five families cold, held-out at fresh +sizes: crop1 **1.0**, flip_h_crop1 **1.0** (the real-travel case, still discovered at the halved +per-start budget), subsample2 **1.0**, upscale2 **0.98**, tile2 **1.0**. Controls: tile2 through the +PLAIN gather **0.14** (the toroidal wrap is load-bearing); no-shape-write **0.0** (unchanged). The +experiment battery's final config: upscale2 1.0/1.0 (M≈I, |t|,|trel| ≤ 0.02), tile2 1.0/1.0 +(M≈2I, trel = 0.4998 — the settle phase centres to ~3 decimal places). + +Synth: `upscale2`/`tile2` added to `SHAPE_TRANSFORMS`. The doubling families' test dims are [3,6] +(outputs ≤ 12×12) to keep the full-budget fit cheap; identifiability (≥2 distinct sizes) unchanged. + +**Deferred (documented):** non-uniform factors (upscale3, tile3 — expected free: seed B's +trel = (k−1)/2 is exact for every k, `1 ≡ 0 (mod 1)` for k=3); mirror-tilings (sign flips near seed +B); wiring `arc_solve --report` (rung b — the user wants BOTH corpus splits re-measured when it +lands); colour on top of shape (rung c). + +## 2026-07-04 11:54 — Rung (b): the shape seam wired into arc_solve; real ARC-AGI-2 v3 measure + +`arc_solve.mojo` now DISPATCHES per task on a closed-form observable of the demos: any train pair +whose dims differ → `ShapeGeomComposedMemory` via `fit_shape_geom` (the rung-(a) two-frame +multi-start); all-same-dims → the unchanged `GeomColorComposedMemory` path (byte-identical for the +same-shape 68%). This is driver-level routing, not a runtime memory-selector — the same-shape +memory PROVABLY scores 0 on the dispatched class, so nothing is being "chosen" that the data +doesn't force. Shape-path scoring: the memory predicts its own output dims from the written rule; +a predicted/true dims mismatch scores that pair 0 (never applied — no OOB). The old +skip-if-every-test-shape-changes shortcut survives only for same-shape-dispatched tasks (where it +remains exact). Each per-task line now carries a trailing `mem: same|shape` marker (appended after +the existing fields — `eval_parallel.sh` reads held-out positionally), giving corpus breakdowns for +free. CI: the full tier's arc_solve leg adds one crop1 shape bundle (fast gate unchanged, ~2 min). +Smoke proof: mixed synth dir {flip_h, crop1, upscale2} → 3/3 solved, markers correct. + +**Eval split v3** (`eval_parallel.sh data_bin/arc2_eval scratch/arc2_eval_v3.txt 10 64 1500`, +10 workers, 11197s): **0 / 120 solved, mean held-out 0.404** (v2: 0.388). Breakdown from the +markers: 39/120 tasks dispatched shape (32.5% — matches corpus_stats exactly); their mean held-out +is **0.054, 0 solved** — the honest first number on the previously-untouchable slice. The +like-for-like control holds: the 81 same-shape tasks mean 0.572 vs v2's implied 0.575 (seeded-ES +noise) — the wiring changed nothing on the same-shape path. Top shape near-misses: 136b0064 (0.66, +gap 0.03) and eee78d87 (0.61) — right predicted dims, partial content. The verdict matches the +rung's documented limits: real eval-split shape tasks overwhelmingly have CONTENT-dependent output +sizes (outside the affine-in-dims rule) or need colour on top of shape (rung c) — the seam is +measured, the expressiveness gaps are now named and quantified. + +**Train split v3** (same harness/budget, 10 workers, 38808s ≈ 10.8h): **22 / 1000 solved (2.2%), +mean held-out 0.501** — more than DOUBLE v2's 10/1000. The marker breakdown decomposes the gain +exactly: +- **shape-dispatched: 320/1000 (32%), 9 solved, mean 0.239** — the first real-ARC shape-changing + solves ever (v2 scored this whole slice 0 by construction): 2dee498d, 60c09cac, 68b67ca3, + 8597cfd7, 8d5021e8, 963e52fc, a416b8f3, be03b35f, c59eb873. Near-misses at 0.89 (53b68214, + 2dc579da). Train's shape slice scores far above eval's (0.239 vs 0.054) — smaller grids, more + affine-in-dims size rules. +- **same-shape: 680/1000, 13 solved, mean 0.625** — net +3 vs v2's 10. Not from this rung: the + same-shape path is byte-identical here; the delta is the **few-demo hardening** (0e138cf), which + landed AFTER the v2 measure and is corpus-measured for the first time now: +5 new solves + (3c9b0459, 5582e5ca, 74dd1130, 9dfd6313, and its designed exhibit **d511f180 — solved at corpus + budget, as the block predicted**), −2 lost (b1948b0a 0.94, ed36ccf7 0.78 — the identity-on-tie + convention's documented trade). + +Raw dumps: `scratch/arc2_eval_v3.txt` / `scratch/arc2_train_v3.txt` (gitignored; per-task +reproducible via SOLVE_SEED). The v3 verdict for the roadmap: the shape seam is now measurable and +productive on the real corpus (train), and the eval split's shape slice names rung (c) — colour on +top of shape — plus content-dependent output sizes as the binding constraints there. + +## 2026-07-04 23:40 — v3 diagnostic breakdown → the expressiveness rung plan (ROADMAP "Next" rewritten) + +Mined the v3 dumps (`scratch/arc2_{train,eval}_v3.txt`, per-task lines with `mem:` markers; all +numbers below reproducible with awk over held-out = field 4, train-fit = field 6) to rank the next +expressiveness rungs on evidence rather than intuition: + +**Train same-shape (680):** 13 solved | **88 near-misses at held-out 0.90–0.99, mean train-fit +0.93** (fail on a FEW cells — the biggest shallow-headroom pool) | 238 at 0.7–0.9 | 146 at <0.4 +with mean train-fit **0.34** — the deep floor: these can't even fit their demos (multi-step / +object-level rules; CMS-chain territory). +**Train shape (320):** 9 solved | quadrants: only 3 tasks are train-high/held-low (what the memory +expresses, it generalizes — again) | **107 with train-fit ≥0.5** (dims + most content fit — +convertible by better content, i.e. colour-on-shape) | **63 with train-fit exactly 0.0** — the +affine-in-dims rule fits NO demo: content-dependent output sizes. +**Eval shape (39):** 19/39 in that dims-never-fit class, 14 partial, 6 with train-fit ≥0.5 — +content-dependent sizes DOMINATE eval's shape slice. + +The ROADMAP "Next" section is rewritten as the evidence-ranked rung ladder, each with its +research/implementation split and named blocker: **C** colour-on-shape (kernel: count signatures +aren't conserved under shape change; validate area-ratio normalization, measure crop's border-loss +robustness at n=3; fallback = correspondence write after a geometry prefit), **S** shape-from- +content (shape WRITE over a small content-statistic basis, residual-selected — bbox-crop the +headline class; audit the 63 ids first), **A** the near-miss audit (measure-first; leading +candidate a self-written mask/gate), **D** k=3/mirror-tiling cheap extensions, **CMS** the depth +chain (expect a wall past depth 2 — literature pass planned at it), with the **GPU gate** as an +explicit zero-capability infrastructure block scheduled immediately before CMS: CPU suffices +through C/S/A/D; the blocker is the `mojo==1.0.0b2` pin (no `gpu` package — MAX migration, all +proof numbers re-proven), not kernel design (the ES is embarrassingly parallel; ~10–30×/fit +realistic). + +Overall split: roughly half research, but each rung's research kernel is ONE identifiable question +— the pattern of the last three landed blocks. No engine code in this entry; next session starts +Rung C. + +## 2026-07-05 — Rung C: colour on top of shape (ShapeGeomColorComposedMemory) + +The top-ranked expressiveness rung. The shape path (`ShapeGeomComposedMemory`) expressed +shape+geometry but had NO colour remapping, so the ~107 convertible train-shape tasks (train-fit +≥0.5, "dims + most content fit") plus eval analogues lost the recolored cells. Rung C composes a +written colour table V on top — the block-5 recipe's THIRD application (after GeomColor, GeomCount): +`out = shape_geom_gather(V(in))`. Colour is cellwise, so it commutes with the copy gather +(colour-then-gather); V is written closed-form, then the unchanged two-frame `fit_shape_geom` runs +on V-pre-mapped demos (the geometry search never sees V — exactly how `fit_geomcolor` reuses +`fit_operator[AttnGatherMemory]`). + +**New surface (all additive; the same-shape core and every existing memory untouched):** +`ShapeGeomColorComposedMemory` (a thin ShapeMemory wrapper: prefix `[0:SHAPEGEOM_DIM]` the +shape+geometry state, suffix `[+COLOR_DIM]` the written V; `apply` = the toroidal shape gather then a +hard V-lookup; `fill_scale` freezes V), `write_color_shaped`, and the `fit_shape_color` driver. +`arc_solve` routes the shape-dispatch branch through it; V=identity makes it byte-identical to the +old shape path. Synth `recolor_{crop1,subsample2,upscale2,tile2}` families + `tests/test_shape_color`. + +**Research kernel #1 — the count-signature write breaks under shape change.** `write_color` +matches per-colour COUNT vectors, assuming count CONSERVATION; shape change breaks it (upscale/tile +multiply every count by the area ratio kr·kc EXACTLY; crop/subsample scale it approximately — +crop drops a colour-dependent border). Fix: normalize each demo's histogram to FRACTIONS (÷ cell +total) before the mismatch — scale-invariant, so `frac_in[c]` matches `frac_out[V(c)]` under any +proportion-preserving resize. Exact for upscale/tile; robust for crop/subsample. + +**Research kernel #2 — the write needs colour-count CONTRAST (measured).** The fraction match +identifies a permutation only when colours have DISTINGUISHABLE frequencies. Under EXACT same-shape +conservation (GeomColor) even a uniform-random grid's tiny fluctuations match exactly; shape change +is LOSSY, so the signal must be REAL contrast — which real ARC grids have (background + sparse +objects, very different counts) and uniform-random grids lack. First measured exhibit: +`recolor_crop1` Ckpt A missed 5/10 colours on uniform grids, recovered fully once grids were drawn +from a per-task palette of distinct frequencies (real-ARC-like). The uniform-crop ceiling is kept +as a control (6/10 missed — the precondition is load-bearing, documented not fought). + +**The strict-superset trap (found via the arc_solve smoke, fixed).** A pure-shape crop task on +UNIFORM grids has no recolor, but the greedy-injective assignment overfits the border-loss noise and +writes a SCRAMBLED V — regressing pure-shape crop from held-out 1.0 to 0.17 (diagnostic: +`V=[6 1 2 4 4 5 6 7 8 3]` instead of identity). Since real corpus shape tasks are overwhelmingly +pure-shape, this would have silently regressed the v3 shape slice. Fix, MEASURED not guessed: +keep `write_color`'s complete greedy-injective matching (full recolor recovery) but add a GLOBAL +ACCEPTANCE GATE — accept the whole map only if its total fraction-mismatch `R_assign < 0.4·R_id` +(the identity assignment's total). A forced injective permutation on a task with NO recolor has a +HIGH residual (measured greedy R_assign/R_id **0.82–1.35** on uniform pure crop across 7 seeds — +forcing a permutation *costs*); a genuine recolor's true permutation has a LOW one (**0.15** +palette crop, **~0** exact-conservation upscale/tile). The 0.4 gate sits wide in that gap +(insensitive across [0.3, 0.55]; real-ARC contrast pushes recolor toward 0). A GLOBAL ratio, not +per-colour — noise averages out (a per-colour 0.5 gate failed: min-over-targets is systematically +below identity on pure noise; and a MUTUAL-best assignment protected pure-shape but under-recovered +recolor, 0.875 — the greedy complete matching is what recovers all colours). `test_shape_color` +guards this: a pure UNIFORM crop1 must keep V=identity AND held-out ≥0.95. + +**Also measured:** subsample-recolor needs adequate output cells (a 2×2 subsampled output starves +the low-frequency colours) — synth subsample dims lifted to {8,10,12}. Real-ARC subsample tasks +aren't 4×4; the identifiability precondition, not a fudge. + +**Result — `test_shape_color` (full tier, cold, held-out at FRESH sizes):** +`recolor_{crop1,subsample2,upscale2,tile2}` each **1.0**; colour-ablation (V=identity) **0.0** (the +colour module is load-bearing); pure tile2 through the colour memory **1.0** and pure uniform crop1 +**1.0 with V=identity** (the strict superset); few-demo n=3 (corpus median) **0.998**; ceiling +control (uniform crop, no contrast) 6/10 colours missed. arc_solve smoke {flip_h, crop1, +recolor_crop1}: flip_h 1.0 (same-shape, byte-identical), crop1 1.0 (pure shape preserved), +recolor_crop1 0.28 (synth uses uniform grids = the documented no-contrast ceiling). CI's full-tier +arc_solve leg adds a `recolor_crop1` bundle so the `fit_shape_color` path runs end-to-end. + +The composition pattern lands a third time; the fraction-signature write + the measured contrast +precondition + the global recolor gate are the new, reusable pieces. Corpus v4 re-measure +(both splits, documented budget) is the separate overnight trigger, deferred by scope. diff --git a/docs/ROADMAP.md b/docs/ROADMAP.md index a085409..00a7183 100644 --- a/docs/ROADMAP.md +++ b/docs/ROADMAP.md @@ -44,9 +44,10 @@ learns both *what* to think and, eventually, *how* to learn. - **Vision A — ARC-AGI 2 (current, active).** Raw held-out solve rate on the real ARC-AGI-2 corpus (`arc_solve --report`) via composable, emergent, self-modifying memories fit in-context by ES — - currently **10/1000 (train) / 0/120 (eval)** at the v2 emergent-composed-memory measure (the M8 - operator ceiling was 5/1000). See "Next — the path to full ARC-AGI 2" below - for the measurable rungs. Still hands the engine *goals* (a task's demonstration pairs are + currently **22/1000 (train) / 0/120 (eval)** at the v3 measure (shape seam wired: 9 shape-changing + solves + few-demo hardening's +3; v2 was 10/1000, the M8 operator ceiling 5/1000); Rung C + (colour-on-shape) is synth-proven and wired, its corpus **v4** re-measure the next overnight + trigger. See "Next — the path to full ARC-AGI 2" below for the measurable rungs. Still hands the engine *goals* (a task's demonstration pairs are compressed supervision) even though it never hands it a DSL. - **Vision B — open-ended mastery (WIP — not yet started, no design work done).** Inspired by Random Network Distillation, open-endedness, and unsupervised RL: an agent that masters its @@ -245,6 +246,58 @@ Concise, milestone-level; each links to `docs/JOURNAL.md` for the full narrative **d511f180 recovered to held-out 1.0** at the corpus budget. n=2's remaining failures are measured UNDERDETERMINATION (exact signature ties — unknowable), documented not fought. (JOURNAL 2026-07-03 08:31.) +- **Shape change — the output-size seam + first family (Next #1, in progress).** The engine was + hard-wired same-shape (every `apply` wrote input-shape cells; `fitness` penalized `in_n ≠ out_n`); + output shape was never a represented quantity. New **`ShapeMemory` trait** (parallel to + `SelfModMemory`, so the same-shape core and all existing memories are untouched): the output shape + is INFERRED IN-CONTEXT — a per-axis affine `out = round(k·in + b)` WRITTEN closed-form by + least-squares over the demo dim-pairs — and produced by the proven AttnGather gather, generalized + so its query grid is the OUTPUT grid and it reads the INPUT grid (`apply_shaped`; same-shape path + bit-identical). The shape rule is frozen (`fill_scale`=0) so the ES still fits only the 7 attention + params, on the B3 landscape. `fitness_shape`/`fit_shape`/`fit_shape_geom` are the shape-aware + generic sibling of the ES core. Proof (`test_shape_change`, cold, held-out at a FRESH input size — + demos drawn at varying sizes so the rule is identifiable, not memorized): `{crop1, flip_h_crop1, + subsample2}` each **held-out 1.0** (subsample2 needed `M=2I`, a 2× scale-up, found cold); shape- + ablation control (no write) **0.0** — the inferred shape rule is load-bearing. Synth + `SHAPE_TRANSFORMS` + `generate_shape_task_groups` (varies demo sizes). **Deferred to follow-on + rungs on this same seam:** upscale/tiling (need a floor/modular gather), wiring `arc_solve --report` + to score the real 32%, and colour composition on top of shape. (JOURNAL 2026-07-03 13:12.) +- **Upscale/tiling — the output-growing shape families (Next #1 rung a).** Correction en route: the + affine gather DOES express blocky upscale exactly (`floor(r/s) = round((r−(s−1)/2)/s)`, no ties) — + only tiling's sawtooth `r mod n` is provably outside it. The mechanism: a **toroidal** output-shaped + gather (wrapped source images — substrate, not primitive), an **extent-relative translation** trel + (absorbs tiling's size-dependent phase), and the query **normalized by the written shape slope** + (resize-as-identity: content never re-learns the scale the shape rule knows; tolerances size-free). + The fit, fixed by a measured grid of ~12 failed configs: a k-fold size change has **two canonical + identity frames** (rescaled plane `M=I`; periodic plane `M=kI, trel=(k−1)/2`), so `fit_shape_geom` + runs the same two cold starts per task — each DISCOVER (soft, temperature searched) then SETTLE + (hard-frozen read, sigma held at the sharp landscape's staircase step scale, alpha decayed) — + winner by demo fitness: an honest multi-start, never task staging. Whole family cold, held-out at + fresh sizes: crop1/flip_h_crop1/subsample2/tile2 **1.0**, upscale2 **0.98**; controls: plain + (non-toroidal) gather fails tile2 (0.14), no-shape-write 0.0. (JOURNAL 2026-07-04 07:40.) +- **Rung (b) — shape seam wired into `arc_solve`; real ARC-AGI-2 v3 measure.** Driver-level dispatch + on the demos (any train pair's dims differ → the shape memory; the alternative provably scores 0, + so it is not a memory-selector), shape-path scoring off the predicted dims, `mem:` markers for + free corpus breakdowns. **Train 22/1000 (v2: 10)** — 9 shape-changing solves (a slice that was 0 + by construction) + the few-demo hardening's net +3 (its exhibit d511f180 solved at corpus budget, + as designed; 2 documented tie-convention losses). **Eval 0/120**, mean 0.404 (v2 0.388): the + 39-task shape slice scores 0.054 — content-dependent output sizes and missing colour-on-shape + (rung c) are the named, quantified constraints there. Same-shape subset bit-consistent with v2 + (the wiring is regression-free). (JOURNAL 2026-07-04 11:54.) +- **Rung C — colour on top of shape (`ShapeGeomColorComposedMemory`).** The composition pattern's + THIRD application: a written colour table V on top of the shape+geometry gather + (`out = shape_geom_gather(V(in))`), fit by the unchanged two-frame `fit_shape_geom` on + V-pre-mapped demos. Research kernel: count conservation breaks under shape change, so V is written + from FRACTION signatures (scale-invariant — exact for upscale/tile, robust for crop/subsample), + which need real colour-count CONTRAST (uniform-random grids can't identify a recolor under a lossy + shape change; real ARC grids can — measured ceiling control). A strict-superset trap fixed: a + greedy write scrambles V on pure-shape low-contrast tasks (crop 1.0→0.17), so a MEASURED global + acceptance gate (`R_assign < 0.4·R_id`; measured gap: pure-shape 0.82–1.35 vs recolor 0.15/≈0) + keeps V=identity unless a recolor clearly explains the demos. `test_shape_color`: all four + `recolor_{crop1,subsample2,upscale2,tile2}` **held-out 1.0 cold at fresh sizes**, colour ablation + 0.0, pure-shape (incl. uniform crop) V=identity 1.0, few-demo n=3 0.998. `arc_solve` shape path + routes through it (byte-identical for pure shape). **Corpus v4 re-measure deferred** (the separate + overnight trigger). (JOURNAL 2026-07-05.) ## Next — the path to full ARC-AGI 2 (Vision A) @@ -252,19 +305,58 @@ Each is its own block, held to the **cold-fit bar** (a scaffolded pass is a nega emergent memories are each measured on the subset they express; the north-star metric is the raw held-out ARC-AGI-2 number. -**Corpus funnel evidence** (2026-07-02, `tools/corpus_stats.py` — the facts the ordering below rests -on): **68% of both splits are same-shape** (680/1000 train, 81/120 eval), so expressiveness — not -shape — is the first binding constraint; median max-grid is **196 cells (train) / 525 (eval)** (ARC -max 900), so real-grid scale is a *compute* constraint (hence the windowed gather and documented -corpus fit-budgets); and **median 3 demos per task (min 2)** vs the synth suite's 8, so every -in-context write rule must hold at 2–3 demonstrations. - -1. **Shape change.** Handle outputs whose dims ≠ inputs — a Domain / output-size generalization - (the output shape itself must be *inferred in-context* from the demos, like any other rule - parameter — never a hand-coded size heuristic). Unlocks the excluded 32% of both splits. -2. **Multi-block CMS chain** (NL §7). Stack memories at multiple update frequencies for multi-step / - object-level reasoning — the (now twice-proven) composition pattern chained in depth, not just in pairs. -3. **Persistent slow weights + task-stream (continual meta-learning).** Stop re-seeding cold per +**Corpus funnel evidence** (2026-07-02, `tools/corpus_stats.py` — plus the 2026-07-05 v3 diagnostic +breakdown, JOURNAL): **68% of both splits are same-shape** (680/1000 train, 81/120 eval); median +max-grid is **196 cells (train) / 525 (eval)** (ARC max 900), so real-grid scale is a *compute* +constraint (hence the windowed gather and documented corpus fit-budgets); **median 3 demos per task +(min 2)** vs the synth suite's 8, so every in-context write rule must hold at 2–3 demonstrations. +The v3 marker breakdown grounds the rung ranking below: train same-shape has **88 near-misses at +held-out 0.90–0.99** (train-fit 0.93 — a few wrong cells) and **146 tasks at a deep floor** +(held-out <0.4, train-fit 0.34 — can't even fit the demos); train shape has **107 tasks with +train-fit ≥0.5** (convertible) and **63 where the affine dims rule fits NO demo**; eval's shape +slice is dominated by that last class (19/39). + +1. **Rung C — colour on top of shape — DONE (synth-proven + wired; corpus v4 deferred).** Landed as + `ShapeGeomColorComposedMemory` (see "Status — done" above). The count-conservation kernel resolved + via scale-invariant FRACTION signatures (contrast-preconditioned) + a measured global recolor + acceptance gate that preserves the pure-shape strict superset. Remaining: the **corpus v4 + re-measure** (both splits, documented budget — the separate overnight trigger) to book the shape- + slice gain on the ~107 convertible train tasks + eval analogues; and the crop-border-loss + correspondence-write fallback stays documented (not needed — fractions + contrast sufficed). +2. **Rung S — shape-from-content** (~40% research / 60% implementation). The 63 train + 19 eval + tasks where NO affine-in-dims rule fits any demo: generalize the shape WRITE (not the gather) to + an affine rule over a small basis of per-demo content statistics — input dims, + non-background bounding-box dims (bbox-crop is a major ARC class), distinct-colour count — basis + selected by least-squares residual across demos (precedent: GeomCount's P-by-residual, block 4's + scoring salience; statistics are representational substrate, not task primitives). Research + part: the emergence-bar argument for the basis, identifiability at 2–3 demos (documented + underdetermination ceiling), background-colour inference. Audit the 63 task ids before coding. +3. **Rung A — the same-shape near-miss audit** (measure-first). The 88 tasks at held-out 0.90–0.99 + fail on a FEW cells: build a tools/ diagnostic that dumps predicted-vs-truth cell diffs for + those ids, clustered by pattern (global colour error vs localized region vs border), and let the + audit name the mechanism — leading candidate: a self-written content MASK/GATE (rule where mask, + identity elsewhere), also the first step toward content-gated composition. Do not design before + the audit; the long-tail hazard (fragmentation into many small classes) is what the audit + decides. +4. **Rung D — small paid-for extensions** (pure implementation). k=3 factors and mirror-tilings + (near the periodic seed B on the two-frame seam); trivial finds from the audit. +5. **Rung CMS — multi-block chain** (NL §7; mostly research). The deep floor (146 train same-shape + tasks at train-fit 0.34, most of eval): multi-step / object-level rules no single memory or pair + expresses — the (twice-proven) composition pattern chained in DEPTH (3+ stages, grid-in/grid-out + intermediates keep the Domain seam). Open questions: per-stage invariant-signal fitting past + depth 2 (no commuting factorization exists for arbitrary triples — expect a block-5-style wall + and plan the literature pass at it), and capacity control (the Schug guardrail). + **GPU gate (infrastructure block, scheduled immediately BEFORE this rung):** CPU is sufficient + through rungs C/S/A/D (synth proofs in minutes; corpus runs overnight at documented budgets) — + the CMS chain and the task-stream (#6) are the 10–100× compute jump. The ES inner loop is + embarrassingly parallel (2N samples × demos × cells × window per iteration; iterations + sequential, so per-launch latency is the kernel-design risk; realistic 10–30× per fit plus + task-level sharding). The REAL blocker is the toolchain: the pinned `mojo==1.0.0b2` slim wheel + has no `gpu` package — GPU means the MAX-platform migration off the hard pin, mechanical but + risky (API churn everywhere; proof numbers must be RE-PROVEN — bit-identity will not survive). + Do it as its own zero-new-capability block with full suite re-verification, never mid-research- + rung (a moving toolchain confounds negative results). +6. **Persistent slow weights + task-stream (continual meta-learning).** Stop re-seeding cold per task: the engine processes a **stream** of tasks and its slow weights **persist**, Reptile-nudged after each in-context fit (M9's outer loop made continual — the prior is never reset). Measurable, in order of strength: (a) at a fixed *narrow* eval budget, solve rate / fit speed **improves with @@ -274,7 +366,7 @@ in-context write rule must hold at 2–3 demonstrations. prior). Known hazard to design around (the M9 lesson: priors help within a *family*): a single flat prior across a heterogeneous stream washes out — the fix is per-family structure that is itself emergent (the Schug hypernetwork route, RESEARCH-NOTES #2: per-task code × shared - templates) and/or the CMS frequency hierarchy (#2), where slow blocks consolidate what fast + templates) and/or the CMS frequency hierarchy (#5), where slow blocks consolidate what fast blocks keep re-discovering. **Serves both visions**: on Vision A it is the meta-learned prior at corpus scale; it is also the tabled **first rung of Vision B** — the same persistence machinery, later driven by self-generated novelty instead of demonstration pairs. diff --git a/run_tests.sh b/run_tests.sh index e225718..43e6406 100755 --- a/run_tests.sh +++ b/run_tests.sh @@ -57,12 +57,19 @@ mojo run -I src src/main.mojo echo "Running held-out generalization driver (src/arc_solve.mojo)..." GEN_DIR="$(mktemp -d)" trap 'rm -rf "$GEN_DIR"' EXIT -python - "$GEN_DIR" <<'PY' +# The full tier adds SHAPE-CHANGING bundles so the driver's shape dispatch +# (ShapeGeomColorComposedMemory + fit_shape_color) runs end-to-end in CI: one +# pure-shape (crop1) and one colour-on-shape (recolor_crop1, Rung C) bundle. The +# fast gate keeps the same-shape-only leg (~a minute cheaper). +python - "$GEN_DIR" "$TIER" <<'PY' import sys sys.path.insert(0, "tools") -from synth_tasks import generate_task_groups +from synth_tasks import generate_task_groups, generate_shape_task_groups generate_task_groups("flip_h", sys.argv[1], num_tasks=2, n_train=6, rows=4, cols=4, seed=0) generate_task_groups("recolor", sys.argv[1], num_tasks=1, n_train=6, rows=4, cols=4, seed=1) +if sys.argv[2] == "full": + generate_shape_task_groups("crop1", sys.argv[1], num_tasks=1, n_train=6, seed=2) + generate_shape_task_groups("recolor_crop1", sys.argv[1], num_tasks=1, n_train=6, seed=3) print("Generated task bundles in", sys.argv[1]) PY mojo run -I src src/arc_solve.mojo "$GEN_DIR"/*.task diff --git a/src/arc_solve.mojo b/src/arc_solve.mojo index e62563b..1ff1ce2 100644 --- a/src/arc_solve.mojo +++ b/src/arc_solve.mojo @@ -2,9 +2,15 @@ from std.sys import argv from std.memory import alloc, UnsafePointer from std.random import seed -from memory_composed import GeomColorComposedMemory, GEOMCOLOR_DIM +from memory_composed import ( + GeomColorComposedMemory, + GEOMCOLOR_DIM, + ShapeGeomColorComposedMemory, + SHAPEGEOMCOLOR_DIM, +) from esper_evolution import ( fit_geomcolor, + fit_shape_color, FIT_N, FIT_ALPHA0, FIT_ALPHA1, @@ -19,16 +25,28 @@ from arc_io import load_arc_task, exact_match # Esper held-out generalization driver. # # Each argument is a path to a `.task` bundle (produced by -# `synth_tasks.generate_task_groups`, or compiled from real ARC-AGI tasks). For -# every task we fit the EMERGENT composed memory (GeomColorComposedMemory: the -# count-signature colour self-write + the windowed attention-gather geometry -# ES — the block-5 retirement of the structured operator) ON THE TRAIN PAIRS -# ONLY, then score it on the held-out TEST pair(s) it never saw. A task is -# solved iff every test pair matches above SOLVE_THRESHOLD. Held-out +# `synth_tasks.generate_task_groups` / `generate_shape_task_groups`, or +# compiled from real ARC-AGI tasks). For every task we fit an EMERGENT +# composed memory ON THE TRAIN PAIRS ONLY, then score it on the held-out TEST +# pair(s) it never saw. Which memory is a DRIVER-LEVEL dispatch on a +# closed-form observable of the demos (NOT a runtime memory-selector — the +# other memory provably scores 0 on the dispatched class): +# - all demo dims equal → GeomColorComposedMemory (block 5: count-signature +# colour self-write + windowed attention-gather geometry ES), +# - any demo changes dims → ShapeGeomColorComposedMemory (the shape seam: per-axis +# affine shape rule written closed-form from the demo dims + the toroidal +# normalized-query gather, two-frame multi-start fit, composed with a written +# colour table V — Rung C; V = identity makes it byte-identical to the pure +# shape+geometry path). Known, documented limits on the real corpus: +# the shape rule is affine in the INPUT DIMS — tasks whose output size +# depends on grid CONTENT mispredict and honestly score 0. +# A task is solved iff every test pair matches above SOLVE_THRESHOLD. Held-out # generalization is uncheatable by memorization. We also report the train-fit # vs held-out gap: `train ~ 1, held-out ~ 0` is memorize-not-generalize; # `train ~ 0` is an expressiveness gap — the breakdown that prioritizes the -# roadmap. +# roadmap. Each per-task line carries a trailing `mem:` marker (same | shape) +# so corpus runs break down by dispatch for free (appended AFTER the existing +# fields — eval_parallel.sh reads held-out positionally as field 4). # # Flags (must precede the task paths): # --report honest-eval mode: 0 solved is a legitimate number, no raise. @@ -63,18 +81,35 @@ def solve_task(task_path: String, n_fit: Int, iters: Int) raises -> Float32: # Deterministic per-task RNG (order/shard invariant) — see SOLVE_SEED. seed(SOLVE_SEED) - # If EVERY test pair is shape-changing, held-out is 0 by construction (the - # memory is same-shape), so the fit cannot change the result — skip it. - # Exact w.r.t. the solve metric; on the real corpus (~32% such tasks) this - # saves a third of the compute. corpus_stats.py filters these out of the - # train-fit diagnostics via the bundle headers. - var any_test_same_shape = False - for i in range(len(task.test)): - if task.test[i].output_grid.size() == task.test[i].input_grid.size(): - any_test_same_shape = True - if not any_test_same_shape: - print(" task:", task_path, " held-out: 0.0 train: 0.0 gap: 0.0") - return 0.0 + # Dispatch on the demos (see the header): any train pair whose DIMS differ + # routes to the shape memory. + var shape_task = False + for i in range(len(task.train)): + if ( + task.train[i].input_grid.rows != task.train[i].output_grid.rows + or task.train[i].input_grid.cols != task.train[i].output_grid.cols + ): + shape_task = True + + # Same-shape dispatch + EVERY test pair shape-changing ⇒ held-out is 0 by + # construction (the same-shape memory cannot express any test pair), so + # the fit cannot change the result — skip it. Exact w.r.t. the solve + # metric. (Shape-dispatched tasks are always fit.) + if not shape_task: + var any_test_same_shape = False + for i in range(len(task.test)): + if ( + task.test[i].output_grid.size() + == task.test[i].input_grid.size() + ): + any_test_same_shape = True + if not any_test_same_shape: + print( + " task:", + task_path, + " held-out: 0.0 train: 0.0 gap: 0.0 mem: same", + ) + return 0.0 # Forward scratch must hold the largest grid the memory touches. var capacity = 1 @@ -89,56 +124,121 @@ def solve_task(task_path: String, n_fit: Int, iters: Int) raises -> Float32: if task.test[i].output_grid.size() > capacity: capacity = task.test[i].output_grid.size() - # Cold per-task fit of the composed memory: colour table written from the - # demos, then the annealed geometry ES on the V-pre-mapped demos. - var state = alloc[Float32](GEOMCOLOR_DIM) - GeomColorComposedMemory.seed(state) - fit_geomcolor( - state, - task.train, - capacity, - n_fit, - FIT_ALPHA0, - FIT_ALPHA1, - FIT_SIGMA0, - FIT_SIGMA1, - iters, - FIT_REG, - ) + # Cold per-task fit of the dispatched composed memory. + var state_dim = GEOMCOLOR_DIM + if shape_task: + state_dim = SHAPEGEOMCOLOR_DIM + var state = alloc[Float32](state_dim) + if shape_task: + # Shape rule + colour table written closed-form from the demos, then the + # two-frame multi-start geometry fit on the V-pre-mapped output-shaped + # gather (Rung C). + ShapeGeomColorComposedMemory.seed(state) + fit_shape_color( + state, + task.train, + capacity, + n_fit, + FIT_ALPHA0, + FIT_ALPHA1, + FIT_SIGMA0, + FIT_SIGMA1, + iters, + FIT_REG, + ) + else: + # Colour table written from the demos, then the annealed geometry ES + # on the V-pre-mapped demos. + GeomColorComposedMemory.seed(state) + fit_geomcolor( + state, + task.train, + capacity, + n_fit, + FIT_ALPHA0, + FIT_ALPHA1, + FIT_SIGMA0, + FIT_SIGMA1, + iters, + FIT_REG, + ) var pred = alloc[Float32](capacity) - # Held-out: minimum exact-match over all test pairs. The memory is - # same-shape, so a test pair whose output dims differ from its input dims is - # inexpressible — it honestly scores 0 (and skipping the compare avoids an - # out-of-bounds read across the mismatched buffers). Real ARC-AGI tasks are - # often shape-changing; this is where the honest number stays low. + # Held-out: minimum exact-match over all test pairs. + # - shape path: the memory PREDICTS its output dims from the written rule; + # a predicted/true dims mismatch scores 0 for that pair (never applied — + # no OOB). This is where content-dependent output sizes honestly fail. + # - same-shape path: a test pair whose output dims differ from its input + # dims is inexpressible — it honestly scores 0 (and skipping the compare + # avoids an out-of-bounds read across the mismatched buffers). var held_out = Float32(1.0) for i in range(len(task.test)): - var rows = task.test[i].input_grid.rows - var cols = task.test[i].input_grid.cols - if task.test[i].output_grid.size() != rows * cols: - held_out = 0.0 - continue - GeomColorComposedMemory.apply(state, task.test[i].input_grid, pred) - var m = exact_match(pred, task.test[i].output_grid.data, rows * cols) + var m = Float32(0.0) + if shape_task: + var pr = ShapeGeomColorComposedMemory.out_rows( + state, task.test[i].input_grid + ) + var pc = ShapeGeomColorComposedMemory.out_cols( + state, task.test[i].input_grid + ) + if ( + pr == task.test[i].output_grid.rows + and pc == task.test[i].output_grid.cols + ): + ShapeGeomColorComposedMemory.apply( + state, task.test[i].input_grid, pr, pc, pred + ) + m = exact_match(pred, task.test[i].output_grid.data, pr * pc) + else: + var rows = task.test[i].input_grid.rows + var cols = task.test[i].input_grid.cols + if task.test[i].output_grid.size() == rows * cols: + GeomColorComposedMemory.apply( + state, task.test[i].input_grid, pred + ) + m = exact_match( + pred, task.test[i].output_grid.data, rows * cols + ) if m < held_out: held_out = m # Train fit (how well the fitted memory reproduces the demos it saw). Same - # same-shape guard as the held-out scoring above. + # dims guards as the held-out scoring above. var train_sum = Float32(0.0) for i in range(len(task.train)): - var rows = task.train[i].input_grid.rows - var cols = task.train[i].input_grid.cols - if task.train[i].output_grid.size() != rows * cols: - continue - GeomColorComposedMemory.apply(state, task.train[i].input_grid, pred) - train_sum += exact_match( - pred, task.train[i].output_grid.data, rows * cols - ) + if shape_task: + var pr = ShapeGeomColorComposedMemory.out_rows( + state, task.train[i].input_grid + ) + var pc = ShapeGeomColorComposedMemory.out_cols( + state, task.train[i].input_grid + ) + if ( + pr != task.train[i].output_grid.rows + or pc != task.train[i].output_grid.cols + ): + continue + ShapeGeomColorComposedMemory.apply( + state, task.train[i].input_grid, pr, pc, pred + ) + train_sum += exact_match( + pred, task.train[i].output_grid.data, pr * pc + ) + else: + var rows = task.train[i].input_grid.rows + var cols = task.train[i].input_grid.cols + if task.train[i].output_grid.size() != rows * cols: + continue + GeomColorComposedMemory.apply(state, task.train[i].input_grid, pred) + train_sum += exact_match( + pred, task.train[i].output_grid.data, rows * cols + ) var train_fit = train_sum / Float32(len(task.train)) + var mem_name = String("same") + if shape_task: + mem_name = String("shape") print( " task:", task_path, @@ -148,6 +248,8 @@ def solve_task(task_path: String, n_fit: Int, iters: Int) raises -> Float32: train_fit, " gap:", train_fit - held_out, + " mem:", + mem_name, ) pred.free() diff --git a/src/esper_evolution.mojo b/src/esper_evolution.mojo index aa1e308..0d6ddb4 100644 --- a/src/esper_evolution.mojo +++ b/src/esper_evolution.mojo @@ -9,9 +9,22 @@ from std.algorithm import parallelize # Dom is a Domain (the Example type + metrics). It consumes generic ExamplePair/ # Task containers so the same fitness loop serves any domain. The concrete # OperatorMemory + ArcGrid are used only by the grid convenience wrapper below. -from memory import Memory, SelfModMemory -from memory_es import OperatorMemory, AttnGatherMemory, ATTN_DIM -from memory_composed import GeomColorComposedMemory, GeomCountComposedMemory +from memory import Memory, SelfModMemory, ShapeMemory +from memory_es import ( + OperatorMemory, + AttnGatherMemory, + ATTN_DIM, + ATTN_BETA_OFF, +) +from memory_composed import ( + GeomColorComposedMemory, + GeomCountComposedMemory, + ShapeGeomComposedMemory, + ShapeGeomSettleMemory, + ShapeGeomColorComposedMemory, + SHAPEGEOM_DIM, + SHAPEGEOM_TREL_OFF, +) from hope import ExamplePair, Task, ArcTaskPair, HopeNode, ArcGrid # Default in-context fit schedule, shared by forward_with_learning and the solve @@ -772,3 +785,357 @@ def fit_geomcount( reg_lambda, ) slow.free() + + +# ========================================== +# Shape-changing memory fit (Vision A / Next #1 — the output-size seam) +# ========================================== +# The shape-aware analogue of `fitness[M]` for a ShapeMemory. The key departure: +# the memory PREDICTS its own output shape from the input (`out_rows`/`out_cols`, +# read off the written shape rule) and writes that many cells, which are scored +# against the demo output at the OUTPUT area. If the predicted shape's area does +# not match the demo output's, the demo is inexpressible under the current shape +# rule — a heavy penalty (the honest analogue of the same-shape guard in +# `fitness`, and it prevents the OOB a mismatched compare would cause). The L2 +# anchor is over the full param vector; the frozen shape-rule slots contribute a +# constant that cancels in the antithetic F+ - F- (they never move), so only the +# content (attention) slots feel it. +def fitness_shape[ + M: ShapeMemory +]( + state: UnsafePointer[Float32, MutAnyOrigin], + slow: UnsafePointer[Float32, MutAnyOrigin], + demos: List[ExamplePair[M.Dom.Example]], + op_output: UnsafePointer[Float32, MutAnyOrigin], + reg_lambda: Float32, +) -> Float32: + var num = len(demos) + if num == 0: + return 0.0 + + var total = Float32(0.0) + for d in range(num): + var pr = M.out_rows(state, demos[d].input_grid) + var pc = M.out_cols(state, demos[d].input_grid) + var out_n = M.Dom.capacity(demos[d].output_grid) + # Predicted output shape wrong for this demo ⇒ inexpressible here. + if pr * pc != out_n: + total += Float32(-1.0e9) + continue + M.apply(state, demos[d].input_grid, pr, pc, op_output) + total += M.Dom.distance(op_output, demos[d].output_grid, out_n) + total = total / Float32(num) + + var pdim = M.param_dim() + var anchor = Float32(0.0) + for i in range(pdim): + var diff = state[i] - slow[i] + anchor += diff * diff + return total - reg_lambda * anchor / Float32(pdim) + + +# Annealed antithetic ES over a ShapeMemory's params — the shape-aware sibling +# of fit_operator/evolve_fast_weights, self-contained (its own scratch, like +# meta_fit_selfmod) since it drives `fitness_shape` and sizes the forward +# scratch to the OUTPUT capacity. Only the scaled params move: the shape rule is +# written by `M.write` before the call and frozen (its `fill_scale` is 0), so +# this search touches only the content (attention) slots. Determinism mirrors +# evolve_fast_weights: serial epsilons in fixed order, parallel independent +# antithetic evals in disjoint per-sample scratch, serial reduction in sample +# order. +def fit_shape[ + M: ShapeMemory +]( + state: UnsafePointer[Float32, MutAnyOrigin], + slow: UnsafePointer[Float32, MutAnyOrigin], + demos: List[ExamplePair[M.Dom.Example]], + grid_capacity: Int, + N: Int, + alpha0: Float32, + alpha1: Float32, + sigma0: Float32, + sigma1: Float32, + iters: Int, + reg_lambda: Float32, +): + if iters <= 0 or N <= 0: + return + var pdim = M.param_dim() + var remainder = pdim % nelts + var rem_start = pdim - remainder + + var eps_all = alloc[Float32](N * pdim) + var pert_all = alloc[Float32](N * pdim) + var op_all = alloc[Float32](N * grid_capacity) + var coeff = alloc[Float32](N) + var grad = alloc[Float32](pdim) + var scale = alloc[Float32](pdim) + M.fill_scale(scale, pdim) + + var alpha_rate = log(alpha1 / alpha0) / Float32(iters) + var sigma_rate = log(sigma1 / sigma0) / Float32(iters) + + for it in range(iters): + var alpha = alpha0 * exp(alpha_rate * Float32(it)) + var sigma = sigma0 * exp(sigma_rate * Float32(it)) + + # Serial: draw all N epsilon vectors (fixed RNG order -> reproducible). + for s in range(N): + var eps_s = eps_all + s * pdim + for j in range(pdim): + eps_s[j] = Float32(randn_float64(0.0, 1.0)) + + # Parallel: antithetic fitness in disjoint per-sample scratch. + @parameter + def sample(s: Int): + var eps_s = eps_all + s * pdim + var pert = pert_all + s * pdim + var op = op_all + s * grid_capacity + + var pos = SIMD[DType.float32, nelts](sigma) + for j in range(0, pdim - nelts + 1, nelts): + var w_vec = state.load[width=nelts](j) + var seps = eps_s.load[width=nelts](j) * scale.load[width=nelts]( + j + ) + pert.store[width=nelts](j, fma(seps, pos, w_vec)) + if remainder > 0: + for j in range(rem_start, pdim): + pert[j] = fma(eps_s[j] * scale[j], sigma, state[j]) + var f_plus = fitness_shape[M](pert, slow, demos, op, reg_lambda) + + var neg = SIMD[DType.float32, nelts](-sigma) + for j in range(0, pdim - nelts + 1, nelts): + var w_vec = state.load[width=nelts](j) + var seps = eps_s.load[width=nelts](j) * scale.load[width=nelts]( + j + ) + pert.store[width=nelts](j, fma(seps, neg, w_vec)) + if remainder > 0: + for j in range(rem_start, pdim): + pert[j] = fma(eps_s[j] * scale[j], -sigma, state[j]) + var f_minus = fitness_shape[M](pert, slow, demos, op, reg_lambda) + + coeff[s] = f_plus - f_minus + + parallelize[sample](N) + + # Serial reduce + step (same shape as evolve_fast_weights). + memset_zero(grad, pdim) + for s in range(N): + var eps_s = eps_all + s * pdim + var c_vec = SIMD[DType.float32, nelts](coeff[s]) + for j in range(0, pdim - nelts + 1, nelts): + var g = grad.load[width=nelts](j) + grad.store[width=nelts]( + j, fma(eps_s.load[width=nelts](j), c_vec, g) + ) + if remainder > 0: + for j in range(rem_start, pdim): + grad[j] = fma(eps_s[j], coeff[s], grad[j]) + + var fac = alpha / (2.0 * Float32(N) * sigma) + var fac_vec = SIMD[DType.float32, nelts](fac) + for j in range(0, pdim - nelts + 1, nelts): + var w_vec = state.load[width=nelts](j) + var sg = grad.load[width=nelts](j) * scale.load[width=nelts](j) + state.store[width=nelts](j, fma(sg, fac_vec, w_vec)) + if remainder > 0: + for j in range(rem_start, pdim): + state[j] = fma(grad[j] * scale[j], fac, state[j]) + + eps_all.free() + pert_all.free() + op_all.free() + coeff.free() + grad.free() + scale.free() + + +# The shape fit's schedule constants (see the ShapeGeomComposedMemory +# temperature comment for the measured failure grid that forces the design). +# DISCOVER runs the standard wide anneal down to SHAPE_SIGMA_FLOOR — the +# sharp-fitness staircase's step scale; below it the antithetic differences +# are almost always zero and the update is pure noise (measured: sigma → 0.05 +# diverges). SETTLE then holds sigma at the floor at the HARD read +# (SHAPE_BETA_READ: beta = 16, probe-exact for the fractional-offset reads, +# frozen via ShapeGeomSettleMemory) and decays alpha — the exact solution is +# a deep plateau there, and the edge gradient parks the state in its interior +# (measured to centre t/trel to ~0.01). Uniform for every task; never +# per-task staging. +comptime SHAPE_BETA_READ = Float32(4.0) # beta = 16: probe-exact hard read +comptime SHAPE_SIGMA_FLOOR = Float32(0.15) # the staircase step scale +comptime SHAPE_SETTLE_DIV = 4 # settle budget = iters / DIV (discovery keeps rest) +comptime SHAPE_SETTLE_ALPHA0 = Float32(0.01) +comptime SHAPE_SETTLE_ALPHA1 = Float32(0.0005) + + +# The per-task in-context fit for ShapeGeomComposedMemory (the shape-change seam +# driver, parallel to fit_geomcolor): write the shape rule closed-form from the +# demos, then run the annealed shape-aware ES on the geometry slots. Constant- +# compute normalization (iters × FIT_DEMO_REF / n_demos) matches the other +# composed drivers — same total demo-evaluations per task. The written shape- +# rule slots are frozen by fill_scale, so the ES moves only the attention + +# trel params, on the proven B3 landscape but now reading the OUTPUT grid. +# Structure (see the SHAPE_* comment above and the two-frame comment on +# ShapeGeomComposedMemory): the SAME TWO cold starts for every task — seed A +# (resized-plane identity) and seed B (periodic-plane identity, from the +# written slope) — each run DISCOVER (wide soft anneal, temperature searched) +# then SETTLE (hard frozen read, sigma held at the staircase step scale, +# alpha decayed), inside the SAME total budget (each start gets half). The +# winner by demo fitness at the hard read is kept: an honest multi-start — +# selection by the task's own train signal, never a task-specific stage. +def fit_shape_geom( + state: UnsafePointer[Float32, MutAnyOrigin], + demos: List[ArcTaskPair], + grid_capacity: Int, + n_fit: Int, + alpha0: Float32, + alpha1: Float32, + sigma0: Float32, + sigma1: Float32, + iters: Int, + reg_lambda: Float32, +) raises: + ShapeGeomComposedMemory.write(state, demos) + + var n_demos = len(demos) + if n_demos < 1: + n_demos = 1 + # Same total budget as every composed driver; each start gets half. + var iters_scaled = iters * FIT_DEMO_REF // n_demos // 2 + var settle_iters = iters_scaled // SHAPE_SETTLE_DIV + var discover_iters = iters_scaled - settle_iters + + # The staircase-viable sigma endpoint (callers pass the generic smooth- + # landscape FIT_SIGMA1, which is below the step scale — floor it). + var sigma_floor = sigma1 + if sigma_floor < SHAPE_SIGMA_FLOOR: + sigma_floor = SHAPE_SIGMA_FLOOR + + var slow = alloc[Float32](SHAPEGEOM_DIM) + var cand = alloc[Float32](SHAPEGEOM_DIM) + var fit_out = alloc[Float32](grid_capacity) + var best_fitness = Float32(-3.0e38) + + for start in range(2): + # Both starts share the written shape rule; only the content frame + # differs. + copy_weights(cand, state, SHAPEGEOM_DIM) + AttnGatherMemory.seed(cand) # M = I, t = 0, soft temperature + cand[SHAPEGEOM_TREL_OFF + 0] = 0.0 + cand[SHAPEGEOM_TREL_OFF + 1] = 0.0 + if start == 1: + ShapeGeomComposedMemory.seed_periodic(cand) + copy_weights(slow, cand, SHAPEGEOM_DIM) + + # DISCOVER: wide annealed search on the soft landscape. + fit_shape[ShapeGeomComposedMemory]( + cand, + slow, + demos, + grid_capacity, + n_fit, + alpha0, + alpha1, + sigma0, + sigma_floor, + discover_iters, + reg_lambda, + ) + + # SETTLE: hard frozen read, sigma held at the step scale, alpha + # decayed — parks the state inside the exact-solution plateau. + cand[ATTN_BETA_OFF] = SHAPE_BETA_READ + fit_shape[ShapeGeomSettleMemory]( + cand, + slow, + demos, + grid_capacity, + n_fit, + SHAPE_SETTLE_ALPHA0, + SHAPE_SETTLE_ALPHA1, + sigma_floor, + sigma_floor, + settle_iters, + reg_lambda, + ) + + # Keep the better start by demo fitness at the hard read (anchor-free: + # slow = cand zeroes the L2 term — pure data signal). + var f = fitness_shape[ShapeGeomSettleMemory]( + cand, cand, demos, fit_out, reg_lambda + ) + if f > best_fitness: + best_fitness = f + copy_weights(state, cand, SHAPEGEOM_DIM) + + slow.free() + cand.free() + fit_out.free() + + +# The per-task in-context fit for ShapeGeomColorComposedMemory (Rung C — +# colour on top of shape), the shape-seam analogue of fit_geomcolor: +# +# 1. WRITE the shape rule + colour table V closed-form from the demos +# (ShapeGeomColorComposedMemory.write — shape via least-squares, V via the +# fraction-normalized count signatures; both geometry-invariant, one pass). +# 2. PRE-MAP the demo inputs through V (colour-then-gather; V is cellwise and +# the gather positional, so they commute — this puts the geometry search on +# the exact same landscape fit_shape_geom already proves, with no colour +# cliff). Dims are unchanged by V, so the written shape rule is identical. +# 3. GEOMETRY: the unchanged two-frame multi-start fit_shape_geom on the +# SHAPEGEOM prefix of `state` (the layout puts it first, so `state` IS the +# shape+geometry weight vector fit_shape_geom expects; V rides in the +# suffix, frozen). fit_shape_geom re-writes the shape rule on the mapped +# demos — idempotent (same dims → same rule). +# +# The pre-map list is built once per task (not per ES iteration), so the hot +# loop stays allocation-free. A pure-shape task writes V = identity, making this +# BYTE-IDENTICAL to fit_shape_geom (the strict-superset regression guard). +def fit_shape_color( + state: UnsafePointer[Float32, MutAnyOrigin], + demos: List[ArcTaskPair], + grid_capacity: Int, + n_fit: Int, + alpha0: Float32, + alpha1: Float32, + sigma0: Float32, + sigma1: Float32, + iters: Int, + reg_lambda: Float32, +) raises: + # 1. Shape rule + colour self-write (fills the shape slots + V table). + ShapeGeomColorComposedMemory.write(state, demos) + + # 2. Pre-map demo inputs through the written V. + var mapped = List[ArcTaskPair]() + for d in range(len(demos)): + var min_g = ArcGrid(demos[d].input_grid.rows, demos[d].input_grid.cols) + for k in range(min_g.rows * min_g.cols): + min_g.data[k] = ShapeGeomColorComposedMemory._v_lookup( + state, demos[d].input_grid.data[k] + ) + var mout = ArcGrid(demos[d].output_grid.rows, demos[d].output_grid.cols) + memcpy( + dest=mout.data, + src=demos[d].output_grid.data, + count=mout.rows * mout.cols, + ) + mapped.append(ArcTaskPair(min_g^, mout^)) + + # 3. The proven two-frame shape geometry fit on the prefix (V frozen in the + # suffix). fit_shape_geom's constant-compute normalization applies as-is. + fit_shape_geom( + state, + mapped, + grid_capacity, + n_fit, + alpha0, + alpha1, + sigma0, + sigma1, + iters, + reg_lambda, + ) diff --git a/src/memory.mojo b/src/memory.mojo index 9bf69b8..99bd696 100644 --- a/src/memory.mojo +++ b/src/memory.mojo @@ -85,3 +85,72 @@ trait SelfModMemory: dst: UnsafePointer[Float32, MutAnyOrigin], ): ... + + +# ========================================== +# Shape-changing memory trait (Vision A / Next #1 — the output-size seam) +# ========================================== +# Every Memory/SelfModMemory above is SAME-SHAPE: `apply` writes exactly +# `inp` cells and the caller reads back `capacity(inp)` cells. A ShapeMemory +# breaks that coupling — it produces an output whose shape DIFFERS from the +# input, with the output shape itself INFERRED IN-CONTEXT from the demos (a rule +# parameter like any other, never a hand-coded size heuristic). It factors, like +# every composed memory (block 5), into two parts fit on signals invariant to +# each other: +# * a SHAPE RULE written closed-form from the demos (`write`) — position-free +# shape arithmetic (in-dims -> out-dims), read back by `out_rows`/`out_cols`; +# * a CONTENT rule (the ES-fit params in `state`) that fills the predicted +# output grid via a gather from the input (`apply` takes the output dims). +# A distinct trait (parallel to SelfModMemory) keeps the same-shape core and all +# existing memories untouched: the shape-aware fitness/fit driver in +# esper_evolution.mojo is generic `[M: ShapeMemory]`, and the seam is additive. +# `param_dim`/`seed`/`fill_scale` mirror Memory (fill_scale ZEROS the written +# shape-rule slots so the ES never moves them — the GeomColor freeze trick). +trait ShapeMemory: + comptime Dom: Domain + + @staticmethod + def param_dim() -> Int: + ... + + @staticmethod + def seed(state: UnsafePointer[Float32, MutAnyOrigin]): + ... + + @staticmethod + def fill_scale(scale: UnsafePointer[Float32, MutAnyOrigin], n: Int): + ... + + # Write the closed-form factors (the shape rule) from the demos — one + # forward pass, never ES-searched. + @staticmethod + def write( + state: UnsafePointer[Float32, MutAnyOrigin], + demos: List[ExamplePair[Self.Dom.Example]], + ): + ... + + # Predicted output dims for `inp` under the written shape rule. + @staticmethod + def out_rows( + state: UnsafePointer[Float32, MutAnyOrigin], inp: Self.Dom.Example + ) -> Int: + ... + + @staticmethod + def out_cols( + state: UnsafePointer[Float32, MutAnyOrigin], inp: Self.Dom.Example + ) -> Int: + ... + + # Produce an `out_rows * out_cols` output into `dst` (the output-shape-aware + # apply — the caller supplies the predicted output dims). + @staticmethod + def apply( + state: UnsafePointer[Float32, MutAnyOrigin], + inp: Self.Dom.Example, + out_rows: Int, + out_cols: Int, + dst: UnsafePointer[Float32, MutAnyOrigin], + ): + ... diff --git a/src/memory_composed.mojo b/src/memory_composed.mojo index 3da8bf3..e1dc74d 100644 --- a/src/memory_composed.mojo +++ b/src/memory_composed.mojo @@ -3,8 +3,14 @@ from std.math import round, exp from std.collections import List, InlineArray from hope import ArcGrid, ArcTaskPair, COLOR_DIM from arc_io import GridDomain -from memory import Memory -from memory_es import AttnGatherMemory, ATTN_DIM +from memory import Memory, ShapeMemory +from memory_es import ( + AttnGatherMemory, + ATTN_DIM, + ATTN_M_OFF, + ATTN_BETA_OFF, + attn_gather_toroidal, +) # ========================================== # Composed geometry × colour memory (ARC-AGI-2 block 5): retire the operator @@ -379,3 +385,536 @@ struct GeomCountComposedMemory(Memory): mh.free() nh.free() m_try.free() + + +# ========================================== +# Composed SHAPE × geometry memory (Vision A / Next #1 — the shape-change seam) +# ========================================== +# The first ShapeMemory: outputs whose dims DIFFER from the input, with the +# output shape INFERRED IN-CONTEXT (never a hand-coded size heuristic). It lifts +# the composition pattern once more — a shape factor written closed-form from the +# demos, composed with the proven AttnGather content gather: +# +# - SHAPE RULE — a per-axis affine out = round(k*in + b), WRITTEN closed-form by +# least-squares over the demo dim-pairs (`write`). Position-free shape +# arithmetic (like write_color/write_content), so it needs no geometry +# knowledge and runs before any search. Covers crop (k=1, b<0), subsample +# (k=1/s), constant output (k=0). When the demos share ONE input size the +# slope/intercept are UNDERDETERMINED (only their combination at that size is +# pinned) — the honest analogue of the few-demo signature ties; the least- +# squares fallback stays exact at the observed size. Identifying k and b (and +# thus generalizing to an UNSEEN input size) requires >= 2 distinct input +# sizes among the demos — which is exactly what the held-out proof supplies. +# - GEOMETRY — the proven AttnGather read run on the OUTPUT grid, through the +# TOROIDAL output-shaped gather (attn_gather_toroidal — the upscale/tiling +# rung), with the query NORMALIZED by the written shape-rule slope +# (resize-as-identity: the content search never re-learns the scale the +# shape rule already knows). M = I is then EVERY pure resize — centred crop, +# subsample, blocky upscale (floor(r/s) = round((r-(s-1)/2)/s) exactly — +# probe-verified); M = ±perm a flip/transpose within the resize; and +# M = kI with an extent-relative translation trel = 1/2 a k x k tiling (the +# sawtooth `r mod n` is the nearest WRAPPED cell of an affine map; a +# constant t cannot cancel tiling's size-dependent n/2 phase across varying +# demo sizes, so that translation is learned in units of the source extent). +# ES-fit over the 6 attention geometry slots + 2 trel slots (the shape rule +# is frozen: fill_scale zeros its slots, the GeomColor freeze trick; the +# temperature is frozen too — see below). +# +# Layout: [0:7] AttnGather content | [7:9] trel (trel_r, trel_c) | [9:13] shape +# rule (kr, br, kc, bc). Colour composition (a write_color pre-map, which +# commutes cellwise) is a documented next-family extension; this block proves +# the shape seam + geometry incl. the modular (upscale/tiling) families. +comptime SHAPEGEOM_TREL_OFF = ATTN_DIM # trel_r, trel_c +comptime SHAPEGEOM_SHAPE_OFF = ATTN_DIM + 2 # kr, br, kc, bc +comptime SHAPEGEOM_DIM = ATTN_DIM + 6 +# ES step scale for trel (see fill_scale: extent-multiplied, seeded at its +# solution — refinement only). +comptime SHAPEGEOM_TREL_SCALE = Float32(0.2) +# TWO IDENTITY FRAMES, TWO FIT REGIMES (a measured grid of ~10 failed +# configurations fixes this design — JOURNAL). A k-fold size change has TWO +# canonical "identity" content maps, and they cannot share one seed: the +# RESIZED plane (M = I under the normalized query — crop, subsample, blocky +# upscale) and the PERIODIC plane (M = kI, trel = (k-1)/2 — tiling reads the +# input as a torus, corner-aligned). Whichever frame a task lives in, its +# geometry is AT or NEAR that frame's seed; asking one seed to travel to the +# other's solution reliably falls into degenerate basins (measured both +# directions). So the fit driver runs the SAME TWO cold starts for every task +# (both derived from the WRITTEN shape slope — never from the task) and keeps +# the better demo fitness: an honest multi-start, not a selector. +# TEMPERATURE: the resize families read at FRACTIONAL offsets (upscale ±1/4), +# so they need a SHARP final read — but a searched temperature is driven soft +# by the ES (it optimizes the Gaussian-smoothed objective, where a soft read +# is robust to the sampler's own jitter), the soft optimum sits displaced +# from the exact solution by more than the snap tolerance, and freezing or +# force-sharpening it during the WIDE search breaks basin discovery. So each +# start runs DISCOVER (temperature searched from the soft seed — the proven +# smooth landscape) then SETTLE (temperature hard-frozen at SHAPE_BETA_READ +# via the ShapeGeomSettleMemory variant, sigma held at the staircase step +# scale, alpha decayed — parks the state inside the exact-solution plateau). +# Uniform across tasks, never per-task staging. + + +struct ShapeGeomComposedMemory(ShapeMemory): + comptime Dom = GridDomain + + @staticmethod + def param_dim() -> Int: + return SHAPEGEOM_DIM + + @staticmethod + def seed(state: UnsafePointer[Float32, MutAnyOrigin]): + # Identity gather (trel = 0: no extent-relative shift) + identity shape + # rule (out == in): the unfit, unwritten memory is the same-shape + # identity. The read temperature starts at AttnGather's soft seed; the + # fit driver owns its sharpening (see the struct comment). + AttnGatherMemory.seed(state) + state[SHAPEGEOM_TREL_OFF + 0] = 0.0 # trel_r + state[SHAPEGEOM_TREL_OFF + 1] = 0.0 # trel_c + state[SHAPEGEOM_SHAPE_OFF + 0] = 1.0 # kr + state[SHAPEGEOM_SHAPE_OFF + 1] = 0.0 # br + state[SHAPEGEOM_SHAPE_OFF + 2] = 1.0 # kc + state[SHAPEGEOM_SHAPE_OFF + 3] = 0.0 # bc + + @staticmethod + def fill_scale(scale: UnsafePointer[Float32, MutAnyOrigin], n: Int): + # Attention keeps its whole preconditioner (temperature SEARCHED — the + # discover phase runs on the proven soft landscape; the settle variant + # below freezes it). trel gets a SMALL scale: it multiplies the source + # extent, so unit-scale jitter swings reads by half the grid — and + # since the two-frame seeds already place trel AT its solution, it + # only ever needs sub-cell refinement. The shape-rule slots are + # WRITTEN, never searched (scale 0 freezes them on any ES path). + AttnGatherMemory.fill_scale(scale, ATTN_DIM) + scale[SHAPEGEOM_TREL_OFF + 0] = SHAPEGEOM_TREL_SCALE + scale[SHAPEGEOM_TREL_OFF + 1] = SHAPEGEOM_TREL_SCALE + for i in range(4): + scale[SHAPEGEOM_SHAPE_OFF + i] = 0.0 + + # Seed B — the PERIODIC identity frame (see the struct comment): the + # content map that reads the input as a corner-aligned torus, M = kI with + # the wrap phase trel = (k-1)/2 that corner-aligns the centred frames + # (k x k tiling is exactly this map). Reads the WRITTEN shape slopes, so + # it must be called AFTER `write`. + @staticmethod + def seed_periodic(state: UnsafePointer[Float32, MutAnyOrigin]): + var kr = state[SHAPEGEOM_SHAPE_OFF + 0] + var kc = state[SHAPEGEOM_SHAPE_OFF + 2] + state[ATTN_M_OFF + 0] = kr + state[ATTN_M_OFF + 3] = kc + state[SHAPEGEOM_TREL_OFF + 0] = (kr - 1.0) * Float32(0.5) + state[SHAPEGEOM_TREL_OFF + 1] = (kc - 1.0) * Float32(0.5) + + # Least-squares slope/intercept of `out` on `in` over the demos for one + # axis, written into state[k_off]/state[b_off]. `axis == 0` fits rows, + # `axis == 1` fits cols. Falls back to the mean ratio (b = 0) when the input + # dimension does not vary — underdetermined but exact at the observed size + # (see the struct comment). + @staticmethod + def _write_axis( + state: UnsafePointer[Float32, MutAnyOrigin], + demos: List[ArcTaskPair], + axis: Int, + k_off: Int, + b_off: Int, + ): + var nd = len(demos) + var s_in = Float32(0.0) + var s_out = Float32(0.0) + var s_in2 = Float32(0.0) + var s_io = Float32(0.0) + for d in range(nd): + var din = demos[d].input_grid.rows + var dout = demos[d].output_grid.rows + if axis != 0: + din = demos[d].input_grid.cols + dout = demos[d].output_grid.cols + var fi = Float32(din) + var fo = Float32(dout) + s_in += fi + s_out += fo + s_in2 += fi * fi + s_io += fi * fo + var fnd = Float32(nd) + var denom = fnd * s_in2 - s_in * s_in + # denom == 0 <=> every demo shares one input size (variance 0). + if denom > Float32(1.0e-6) or denom < Float32(-1.0e-6): + var k = (fnd * s_io - s_in * s_out) / denom + state[k_off] = k + state[b_off] = (s_out - k * s_in) / fnd + else: + # Fallback: the mean out/in ratio with zero intercept. + var k = Float32(1.0) + if s_in > Float32(0.0): + k = s_out / s_in + state[k_off] = k + state[b_off] = Float32(0.0) + + @staticmethod + def write( + state: UnsafePointer[Float32, MutAnyOrigin], + demos: List[ArcTaskPair], + ): + if len(demos) == 0: + return + Self._write_axis( + state, demos, 0, SHAPEGEOM_SHAPE_OFF + 0, SHAPEGEOM_SHAPE_OFF + 1 + ) + Self._write_axis( + state, demos, 1, SHAPEGEOM_SHAPE_OFF + 2, SHAPEGEOM_SHAPE_OFF + 3 + ) + + @staticmethod + def out_rows( + state: UnsafePointer[Float32, MutAnyOrigin], inp: ArcGrid + ) -> Int: + var v = round( + state[SHAPEGEOM_SHAPE_OFF + 0] * Float32(inp.rows) + + state[SHAPEGEOM_SHAPE_OFF + 1] + ) + var r = Int(v) + if r < 1: + r = 1 + return r + + @staticmethod + def out_cols( + state: UnsafePointer[Float32, MutAnyOrigin], inp: ArcGrid + ) -> Int: + var v = round( + state[SHAPEGEOM_SHAPE_OFF + 2] * Float32(inp.cols) + + state[SHAPEGEOM_SHAPE_OFF + 3] + ) + var c = Int(v) + if c < 1: + c = 1 + return c + + @staticmethod + def apply( + state: UnsafePointer[Float32, MutAnyOrigin], + inp: ArcGrid, + out_rows: Int, + out_cols: Int, + dst: UnsafePointer[Float32, MutAnyOrigin], + ): + # The attention slots [0:7] are `state` itself; trel and the written + # shape-rule slopes ride in explicitly (the slopes NORMALIZE the query: + # resize-as-identity — see attn_gather_toroidal). The source is + # toroidal, which is what makes the modular families (tiling) + # expressible. + attn_gather_toroidal( + state, + state[SHAPEGEOM_TREL_OFF + 0], + state[SHAPEGEOM_TREL_OFF + 1], + state[SHAPEGEOM_SHAPE_OFF + 0], + state[SHAPEGEOM_SHAPE_OFF + 2], + inp, + out_rows, + out_cols, + dst, + ) + + +# The SETTLE-phase variant of ShapeGeomComposedMemory (same layout, same +# gather, same writes — pure delegation) with ONE difference: the temperature +# slot is FROZEN (fill_scale 0). The settle phase runs at the hard read +# (SHAPE_BETA_READ, set by the driver) and a searched temperature would be +# driven soft again by the ES; fill_scale is a static per-type property, so +# the phase difference is a type. See the temperature comment above. +struct ShapeGeomSettleMemory(ShapeMemory): + comptime Dom = GridDomain + + @staticmethod + def param_dim() -> Int: + return SHAPEGEOM_DIM + + @staticmethod + def seed(state: UnsafePointer[Float32, MutAnyOrigin]): + ShapeGeomComposedMemory.seed(state) + + @staticmethod + def fill_scale(scale: UnsafePointer[Float32, MutAnyOrigin], n: Int): + ShapeGeomComposedMemory.fill_scale(scale, n) + scale[ATTN_BETA_OFF] = 0.0 + + @staticmethod + def write( + state: UnsafePointer[Float32, MutAnyOrigin], + demos: List[ArcTaskPair], + ): + ShapeGeomComposedMemory.write(state, demos) + + @staticmethod + def out_rows( + state: UnsafePointer[Float32, MutAnyOrigin], inp: ArcGrid + ) -> Int: + return ShapeGeomComposedMemory.out_rows(state, inp) + + @staticmethod + def out_cols( + state: UnsafePointer[Float32, MutAnyOrigin], inp: ArcGrid + ) -> Int: + return ShapeGeomComposedMemory.out_cols(state, inp) + + @staticmethod + def apply( + state: UnsafePointer[Float32, MutAnyOrigin], + inp: ArcGrid, + out_rows: Int, + out_cols: Int, + dst: UnsafePointer[Float32, MutAnyOrigin], + ): + ShapeGeomComposedMemory.apply(state, inp, out_rows, out_cols, dst) + + +# ========================================== +# Composed SHAPE × geometry × COLOUR memory (Vision A / Next #1, Rung C) +# ========================================== +# Colour on top of the shape seam: the third application of the block-5 +# composition recipe (after GeomColor and GeomCount). The shape path +# (ShapeGeomComposedMemory) expresses shape + geometry but has NO colour +# remapping — any shape-change task that also recolors fails on the recolored +# cells. This wrapper composes a written colour table V on top, bringing the +# shape path to parity with the same-shape GeomColorComposedMemory. +# +# out = shape_geom_gather( V(in) ) +# +# Colour is CELLWISE, so it commutes with the copy gather (mapping a cell's +# colour before or after selecting it is identical) — the same colour-then-gather +# decoupling as GeomColor. V is written closed-form, then the proven two-frame +# shape geometry fit runs on demos PRE-MAPPED through V (fit_shape_color in +# esper_evolution.mojo), so the geometry search is unchanged — it never sees V. +# +# THE RESEARCH KERNEL — the count-signature write under shape change. The +# GeomColor write_color matches per-colour COUNT vectors across demos, which +# assumes count CONSERVATION; shape change breaks it (upscale/tile multiply every +# count by the area ratio kr*kc EXACTLY; crop/subsample scale it approximately — +# crop loses a colour-dependent border). The fix (write_color_shaped): normalize +# each demo's colour histogram to FRACTIONS (÷ cell total) before the mismatch +# matrix. Fractions are SCALE-INVARIANT, so frac_in[c] matches frac_out[V(c)] +# under any proportion-preserving resize — exact for upscale/tile, robust for +# crop/subsample (only border/subsample sampling noise). This is the ROADMAP's +# "normalize by the area ratio kr*kc" framing implemented scale-free (handles +# crop's varying per-demo ratio automatically). Fractions break the integer-exact +# tie test the few-demo hardening relied on, so the tie test carries a small +# tolerance (SHAPEGEOMCOLOR_TIE_TOL); the injective-assignment robustness at +# n=2/3 is measured in test_shape_color (the rung's one research question). +# +# Layout: [0:SHAPEGEOM_DIM] the shape+geometry state (attn | trel | shape rule) +# | [SHAPEGEOM_DIM : +COLOR_DIM] the written colour table V. fill_scale zeroes +# V (written, never searched) and delegates the prefix to ShapeGeom, so the ES +# still moves only the attention + trel slots on the proven B3 landscape. +comptime SHAPEGEOMCOLOR_V_OFF = SHAPEGEOM_DIM # the 10-entry written colour table +comptime SHAPEGEOMCOLOR_DIM = SHAPEGEOM_DIM + COLOR_DIM +# Fraction mismatches are not integer-exact, so genuine assignment ties (two +# colours with identical across-demo fraction signatures) are compared within a +# small tolerance rather than by ==; identity is still preferred on a tie. +comptime SHAPEGEOMCOLOR_TIE_TOL = Float32(1.0e-6) +# GLOBAL recolor-acceptance gate (see write_color_shaped's assignment). The +# whole non-identity colour map is accepted only when the mutual-best +# assignment's total fraction-mismatch is below this fraction of the IDENTITY +# assignment's total — i.e. a recolor is written only when a permutation +# explains the demos clearly better than "no colour change". A genuine recolor +# clears it wide (measured crop-recolor ratio 0.22; exact-conservation upscale/ +# tile ~0); a PURE-shape task on low-contrast/lossy grids does NOT (measured +# uniform-crop pure ratio 0.61-0.99 across seeds), so V defaults to identity and +# the colour path stays a STRICT SUPERSET of the pure-shape path (without the +# gate a noisy uniform crop scrambles V and regresses pure-shape crop 1.0 -> +# 0.17). A GLOBAL ratio (not per-colour) so the noise averages out; the 0.4 +# threshold sits in the measured gap [0.22 recolor | 0.61 pure] and is +# insensitive across [0.3, 0.55]. Real-ARC contrast pushes recolor toward 0. +comptime SHAPEGEOMCOLOR_RECOLOR_GATE = Float32(0.4) + + +struct ShapeGeomColorComposedMemory(ShapeMemory): + comptime Dom = GridDomain + + @staticmethod + def param_dim() -> Int: + return SHAPEGEOMCOLOR_DIM + + @staticmethod + def seed(state: UnsafePointer[Float32, MutAnyOrigin]): + # Identity shape+geometry (same-shape identity) + identity colour table: + # the unfit, unwritten memory is the identity transform. + ShapeGeomComposedMemory.seed(state) + for s in range(COLOR_DIM): + state[SHAPEGEOMCOLOR_V_OFF + s] = Float32(s) + + @staticmethod + def fill_scale(scale: UnsafePointer[Float32, MutAnyOrigin], n: Int): + # Prefix keeps ShapeGeom's whole preconditioner (attn searched, trel + # small, shape rule frozen); the V group gets scale 0 — V is WRITTEN + # from the demos, never searched (perturbation AND update are both + # scale-multiplied, so V never moves on any ES path). + ShapeGeomComposedMemory.fill_scale(scale, SHAPEGEOM_DIM) + for s in range(COLOR_DIM): + scale[SHAPEGEOMCOLOR_V_OFF + s] = 0.0 + + @staticmethod + def _v_lookup( + state: UnsafePointer[Float32, MutAnyOrigin], val: Float32 + ) -> Float32: + # Hard colour read: nearest-integer index into the written table, output + # rounded (the written V entries are ~integer). Mirrors + # GeomColorComposedMemory._v_lookup at the shape layout's V offset. + var idx = Int(round(val)) + if idx < 0: + idx = 0 + if idx > COLOR_DIM - 1: + idx = COLOR_DIM - 1 + return round(state[SHAPEGEOMCOLOR_V_OFF + idx]) + + @staticmethod + def write( + state: UnsafePointer[Float32, MutAnyOrigin], + demos: List[ArcTaskPair], + ): + # Shape rule (least-squares affine over the demo dim-pairs) + the colour + # table V (fraction-normalized count signatures). Both closed-form, one + # pass, never ES-searched; both geometry-invariant, so they run before + # any search. + ShapeGeomComposedMemory.write(state, demos) + Self.write_color_shaped(state, demos) + + @staticmethod + def out_rows( + state: UnsafePointer[Float32, MutAnyOrigin], inp: ArcGrid + ) -> Int: + return ShapeGeomComposedMemory.out_rows(state, inp) + + @staticmethod + def out_cols( + state: UnsafePointer[Float32, MutAnyOrigin], inp: ArcGrid + ) -> Int: + return ShapeGeomComposedMemory.out_cols(state, inp) + + @staticmethod + def apply( + state: UnsafePointer[Float32, MutAnyOrigin], + inp: ArcGrid, + out_rows: Int, + out_cols: Int, + dst: UnsafePointer[Float32, MutAnyOrigin], + ): + # Gather-then-colour (the eval order; equals the colour-then-gather of + # the fit once the read is sharp — every use runs the fit first, same as + # GeomColorComposedMemory.apply). The prefix drives the toroidal + # output-shaped gather, then the hard V read maps each output cell. + ShapeGeomComposedMemory.apply(state, inp, out_rows, out_cols, dst) + for k in range(out_rows * out_cols): + dst[k] = Self._v_lookup(state, dst[k]) + + @staticmethod + def write_color_shaped( + state: UnsafePointer[Float32, MutAnyOrigin], + demos: List[ArcTaskPair], + ): + """Write the colour table V from the demos' FRACTION signatures. + + The shape-invariant sibling of GeomColorComposedMemory.write_color: per + demo the per-colour input/output histograms are normalized to FRACTIONS + (÷ that grid's cell total) before the across-demo squared-difference + mismatch, so the match survives the count rescaling a shape change + induces (exact for area-ratio families, robust for crop/subsample). The + assignment is the same global-min greedy INJECTIVE matching with identity + defaults for unseen colours and identity preference on (tolerance-)ties. + Stack accumulators only — no allocation. + """ + var num = len(demos) + if num == 0: + return + var mismatch = InlineArray[Float32, COLOR_DIM * COLOR_DIM](fill=0.0) + var seen = InlineArray[Int, COLOR_DIM](fill=0) + for d in range(num): + var cnt_in = InlineArray[Float32, COLOR_DIM](fill=0.0) + var cnt_out = InlineArray[Float32, COLOR_DIM](fill=0.0) + var in_n = demos[d].input_grid.rows * demos[d].input_grid.cols + var out_n = demos[d].output_grid.rows * demos[d].output_grid.cols + for k in range(in_n): + var ci = Int(round(demos[d].input_grid.data[k])) + if ci >= 0 and ci < COLOR_DIM: + cnt_in[ci] += 1.0 + seen[ci] = 1 + for k in range(out_n): + var co = Int(round(demos[d].output_grid.data[k])) + if co >= 0 and co < COLOR_DIM: + cnt_out[co] += 1.0 + # Normalize to fractions (÷ cell total) — the scale-invariant fix. + var inv_in = Float32(1.0) / Float32(in_n) if in_n > 0 else Float32( + 0.0 + ) + var inv_out = Float32(1.0) / Float32( + out_n + ) if out_n > 0 else Float32(0.0) + for c in range(COLOR_DIM): + cnt_in[c] *= inv_in + cnt_out[c] *= inv_out + for c in range(COLOR_DIM): + for c2 in range(COLOR_DIM): + var diff = cnt_in[c] - cnt_out[c2] + mismatch[c * COLOR_DIM + c2] += diff * diff + # Assignment: write_color's GLOBAL-MIN GREEDY INJECTIVE matching (the + # few-demo-hardened complete permutation — identity for unseen, identity + # preference on ties within SHAPEGEOMCOLOR_TIE_TOL since fractions aren't + # integer-exact), then a GLOBAL ACCEPTANCE GATE. Greedy alone is safe on + # SAME-shape (exact count conservation ⇒ identity is exactly zero and + # always wins for pure geometry) but a lossy shape change leaves the + # fractions NOISY, so on a pure-shape task greedy still writes SOME + # scrambled permutation (measured: pure crop1 1.0 -> 0.17). The gate + # decides accept-vs-identity for the WHOLE map: a forced injective + # permutation on a task with NO recolor has a HIGH total residual + # (measured greedy R_assign/R_id 0.82-1.35 on uniform pure crop across + # seeds — forcing a permutation costs), whereas a genuine recolor's true + # permutation has a LOW one (0.15 palette crop; ~0 exact-conservation + # upscale/tile). So V = the greedy map iff R_assign < gate*R_id, else + # identity — full recolor recovery AND the pure-shape strict superset. + var target = InlineArray[Int, COLOR_DIM](fill=0) + var assigned = InlineArray[Int, COLOR_DIM](fill=0) + var taken = InlineArray[Int, COLOR_DIM](fill=0) + var n_seen = 0 + for c in range(COLOR_DIM): + target[c] = c # identity default (unseen colours stay here) + if seen[c] == 0: + assigned[c] = 1 + else: + n_seen += 1 + for _ in range(n_seen): + var best_m = Float32(1.0e30) + var best_c = -1 + var best_t = -1 + for c in range(COLOR_DIM): + if assigned[c] == 1: + continue + for t in range(COLOR_DIM): + if taken[t] == 1: + continue + var m = mismatch[c * COLOR_DIM + t] + var better = m < best_m - SHAPEGEOMCOLOR_TIE_TOL + if abs(m - best_m) <= SHAPEGEOMCOLOR_TIE_TOL: + # Tolerance-tie: prefer an identity pair, then lower idx. + var new_id = 1 if c == t else 0 + var cur_id = 1 if best_c == best_t else 0 + if new_id > cur_id: + better = True + if better: + best_m = m + best_c = c + best_t = t + target[best_c] = best_t + assigned[best_c] = 1 + taken[best_t] = 1 + # Global acceptance gate (see SHAPEGEOMCOLOR_RECOLOR_GATE). r_id ~ 0 + # (exact-conservation pure shape) fails the strict `<` ⇒ identity. + var r_id = Float32(0.0) + var r_assign = Float32(0.0) + for c in range(COLOR_DIM): + if seen[c] == 0: + continue + r_id += mismatch[c * COLOR_DIM + c] + r_assign += mismatch[c * COLOR_DIM + target[c]] + var accept = r_assign < SHAPEGEOMCOLOR_RECOLOR_GATE * r_id + for c in range(COLOR_DIM): + state[SHAPEGEOMCOLOR_V_OFF + c] = Float32( + target[c] if accept else c + ) diff --git a/src/memory_es.mojo b/src/memory_es.mojo index ae5bb35..3784b4d 100644 --- a/src/memory_es.mojo +++ b/src/memory_es.mojo @@ -371,6 +371,28 @@ struct AttnGatherMemory(Memory): weights: UnsafePointer[Float32, MutAnyOrigin], inp: ArcGrid, dst: UnsafePointer[Float32, MutAnyOrigin], + ): + # Same-shape path: query grid == input grid. Delegates to the + # output-shape-aware gather with out dims = in dims, so this is + # BIT-IDENTICAL to the pre-shape-seam apply (the query centre, source + # centre, projection and windowed softmax all reduce to the old code). + Self.apply_shaped(weights, inp, inp.rows, inp.cols, dst) + + # Output-shape-aware gather (the shape-change seam, Vision A / Next #1). The + # QUERY grid is (out_rows, out_cols) centred on the OUTPUT extent, while the + # gather still reads the INPUT grid (inp.rows/cols) centred on the INPUT + # extent — the learned projection q = M*v_out + t maps an output coordinate + # into the input's centred frame. Decoupling query size from source size is + # the whole change: M = I, t = 0 reads the centred input (a centred crop), + # M = sI a subsample by s, M = ±perm a flip/transpose within the resize. + # For out == in this is exactly the same-shape gather above. + @staticmethod + def apply_shaped( + weights: UnsafePointer[Float32, MutAnyOrigin], + inp: ArcGrid, + out_rows: Int, + out_cols: Int, + dst: UnsafePointer[Float32, MutAnyOrigin], ): var m00 = weights[ATTN_M_OFF + 0] var m01 = weights[ATTN_M_OFF + 1] @@ -380,15 +402,19 @@ struct AttnGatherMemory(Memory): var t_c = weights[ATTN_T_OFF + 1] var beta_raw = weights[ATTN_BETA_OFF] var beta = beta_raw * beta_raw # >= 0, no exp overflow risk + # Source (gather) grid: the input's extent and centre. var rows = inp.rows var cols = inp.cols var cr = Float32(rows - 1) * Float32(0.5) var cc = Float32(cols - 1) * Float32(0.5) - - for r in range(rows): - var vr = Float32(r) - cr - for c in range(cols): - var vc = Float32(c) - cc + # Query grid: the OUTPUT extent and centre (== input's when out == in). + var cr_out = Float32(out_rows - 1) * Float32(0.5) + var cc_out = Float32(out_cols - 1) * Float32(0.5) + + for r in range(out_rows): + var vr = Float32(r) - cr_out + for c in range(out_cols): + var vc = Float32(c) - cc_out # Read target q = M*v + t (where in coord space to gather from). var qr = fma(m00, vr, fma(m01, vc, t_r)) var qc = fma(m10, vr, fma(m11, vc, t_c)) @@ -452,4 +478,144 @@ struct AttnGatherMemory(Memory): var w = exp(score - max_score) z += w s = fma(w, inp.data[rj * cols + cj], s) - dst[r * cols + c] = s / z + dst[r * out_cols + c] = s / z + + +# TOROIDAL output-shaped gather (the upscale/tiling family, Vision A / Next #1 +# rung a). Same learned attention read as `apply_shaped` (`weights` is the same +# 7-slot attention block) with two changes that make MODULAR content rules +# expressible: +# +# - The source grid is a TORUS: the coordinate displacement per axis is wrapped +# into (-extent/2, extent/2] (`d -= extent*round(d/extent)`), so a query past +# the input's edge reads the input's periodic image — tiling's sawtooth +# `out[r] = in[r mod rows]`, provably outside any single affine (M, t), is +# the nearest wrapped cell of an affine map. Toroidal topology is a substrate +# choice (precedent: the selfmod-grid memories' toroidal neighbourhoods), not +# a task primitive: which modular rule (if any) is read is entirely in the +# learned (M, t, trel, beta). For queries that stay in range (crop / flip / +# subsample / upscale) the wrap touches only the far softmax tail — at any +# sharp temperature the nearest source cell is unchanged. +# - `trel_r`/`trel_c` are EXTENT-RELATIVE translations: q += trel*extent. The +# centred query/source frames leave tiling with a size-dependent phase +# (n/2 for tile-2) that a constant t cannot cancel across the varying demo +# sizes; a translation in units of the source extent absorbs it with one +# size-independent parameter (trel = 1/2 for tile-2), fit by the ES exactly +# like t. +# - `kr`/`kc` NORMALIZE the query by the (written, frozen) shape-rule slope: +# q = M*(v_out/k) + ... — resize-as-identity. The content search must not +# re-learn the scale the shape rule already knows: without this, a resize +# family's M is 1/k and its exactness tolerance shrinks with the output +# extent (measured: upscale-2's m11 needs ±0.045, far below the ES's +# settling noise at the staircase-viable sigma — the fit reliably parked +# just off the plateau). Normalized, M = I (the seed) IS the pure-resize +# solution for every k, and all tolerances are size-free. k <= 0 (a written +# constant-output rule) falls back to 1. +# +# The window scan is the wrapped analogue of `apply_shaped`'s: the span is +# capped at the torus period (scanning more would visit a source cell twice +# through its periodic images), the centre is wrapped into bounds, and indices +# wrap modularly. For extents <= 2*ATTN_WINDOW+1 every source cell is scanned, +# so synth-scale results have no window truncation at all. +def attn_gather_toroidal( + weights: UnsafePointer[Float32, MutAnyOrigin], + trel_r: Float32, + trel_c: Float32, + kr: Float32, + kc: Float32, + inp: ArcGrid, + out_rows: Int, + out_cols: Int, + dst: UnsafePointer[Float32, MutAnyOrigin], +): + var m00 = weights[ATTN_M_OFF + 0] + var m01 = weights[ATTN_M_OFF + 1] + var m10 = weights[ATTN_M_OFF + 2] + var m11 = weights[ATTN_M_OFF + 3] + var t_r = weights[ATTN_T_OFF + 0] + var t_c = weights[ATTN_T_OFF + 1] + var beta_raw = weights[ATTN_BETA_OFF] + var beta = beta_raw * beta_raw + var rows = inp.rows + var cols = inp.cols + var frows = Float32(rows) + var fcols = Float32(cols) + var cr = Float32(rows - 1) * Float32(0.5) + var cc = Float32(cols - 1) * Float32(0.5) + var cr_out = Float32(out_rows - 1) * Float32(0.5) + var cc_out = Float32(out_cols - 1) * Float32(0.5) + + # Query normalization by the shape-rule slope (see the header). + var inv_kr = Float32(1.0) + if kr > Float32(1.0e-3): + inv_kr = 1.0 / kr + var inv_kc = Float32(1.0) + if kc > Float32(1.0e-3): + inv_kc = 1.0 / kc + + # Window span, capped at the torus period per axis (see the header). + var span_r = 2 * ATTN_WINDOW + 1 + if span_r > rows: + span_r = rows + var span_c = 2 * ATTN_WINDOW + 1 + if span_c > cols: + span_c = cols + + for r in range(out_rows): + var vr = (Float32(r) - cr_out) * inv_kr + for c in range(out_cols): + var vc = (Float32(c) - cc_out) * inv_kc + # q = M*v + t + trel*extent, in the input's centred frame. + var qr = fma(m00, vr, fma(m01, vc, t_r)) + trel_r * frows + var qc = fma(m10, vr, fma(m11, vc, t_c)) + trel_c * fcols + + # Window start: centre the span on q's wrapped source index. + var ctr_r = Int(round(qr + cr)) % rows + if ctr_r < 0: + ctr_r += rows + var ctr_c = Int(round(qc + cc)) % cols + if ctr_c < 0: + ctr_c += cols + var r0 = ctr_r - ATTN_WINDOW + var c0 = ctr_c - ATTN_WINDOW + + # Pass 1: max score over the wrapped window (numerical stability). + var max_score = Float32(-1.0e30) + for kr in range(span_r): + var rj = (r0 + kr) % rows + if rj < 0: + rj += rows + var dvr = qr - (Float32(rj) - cr) + dvr -= frows * round(dvr / frows) + var dvr2 = dvr * dvr + for kc in range(span_c): + var cj = (c0 + kc) % cols + if cj < 0: + cj += cols + var dvc = qc - (Float32(cj) - cc) + dvc -= fcols * round(dvc / fcols) + var score = -beta * (dvr2 + dvc * dvc) + if score > max_score: + max_score = score + + # Pass 2: softmax-weighted gather (streaming, no per-cell buffer). + var z = Float32(0.0) + var s = Float32(0.0) + for kr in range(span_r): + var rj = (r0 + kr) % rows + if rj < 0: + rj += rows + var dvr = qr - (Float32(rj) - cr) + dvr -= frows * round(dvr / frows) + var dvr2 = dvr * dvr + for kc in range(span_c): + var cj = (c0 + kc) % cols + if cj < 0: + cj += cols + var dvc = qc - (Float32(cj) - cc) + dvc -= fcols * round(dvc / fcols) + var score = -beta * (dvr2 + dvc * dvc) + var w = exp(score - max_score) + z += w + s = fma(w, inp.data[rj * cols + cj], s) + dst[r * out_cols + c] = s / z diff --git a/tests/test_shape_change.mojo b/tests/test_shape_change.mojo new file mode 100644 index 0000000..e8ce7b1 --- /dev/null +++ b/tests/test_shape_change.mojo @@ -0,0 +1,314 @@ +# suite-tier: full +from std.memory import alloc, UnsafePointer +from std.random import seed, random_float64 +from std.math import round +from std.collections import List + +# Run from the project root: `mojo run -I src tests/test_shape_change.mojo`. +from hope import ArcGrid, ArcTaskPair +from memory_composed import ( + ShapeGeomComposedMemory, + SHAPEGEOM_DIM, + SHAPEGEOM_SHAPE_OFF, +) +from memory_es import AttnGatherMemory, ATTN_BETA_OFF +from esper_evolution import ( + fit_shape_geom, + fit_shape, + FIT_N, + FIT_ALPHA0, + FIT_ALPHA1, + FIT_SIGMA0, + FIT_SIGMA1, + FIT_ITERS, + FIT_REG, + SHAPE_BETA_READ, +) +from arc_io import exact_match + +# ========================================================================== +# The SHAPE-CHANGE seam proof (Vision A / Next #1): ShapeGeomComposedMemory +# produces outputs whose dims DIFFER from the input, with the output shape +# INFERRED IN-CONTEXT from the demos (a closed-form per-axis affine shape rule), +# composed with the proven AttnGather content gather run on the OUTPUT grid. +# +# Each task's demos are drawn at VARYING input sizes, so the shape rule +# out = k*in + b is genuinely identifiable (not memorized), and the held-out +# test — a FRESH, unseen input size — is an uncheatable generalization probe. +# +# Ckpt A — the shape write is exact: on crop1 the least-squares rule recovers +# (k=1, b=-2) per axis from the varying-size demos. +# Ckpt B — the seam + geometry bar: {crop1, flip_h_crop1, subsample2, +# upscale2, tile2} each >= 0.95 held-out at a fresh size, per-task +# cold (seed → one fit_shape_geom call). upscale2 (blocky +# replication, M = I/2 at sharp temperature) and tile2 (the modular +# sawtooth `in[r mod n]` — nearest WRAPPED cell of the affine map at +# trel = 1/2) are the output-GROWING families the toroidal gather +# exists for. +# Control — shape ablation: the SAME content fit WITHOUT the shape write +# (identity shape rule) predicts the wrong output size on every +# demo/test, so held-out collapses to ~0 — the inferred shape rule is +# load-bearing, not scaffolding. +# Control — wrap ablation: the fitted tile2 state read through the PLAIN +# (non-toroidal) gather collapses — the modular source addressing is +# load-bearing for tiling, exactly as the sawtooth argument says. +# ========================================================================== + + +def rand_grid(rows: Int, cols: Int) -> ArcGrid: + var g = ArcGrid(rows, cols) + for k in range(rows * cols): + g.data[k] = Float32(Int(random_float64(0.0, 10.0))) + return g^ + + +# A random axis length per family (rows and cols drawn independently; the same +# RNG stream as the grids, for determinism). subsample2 needs even dims (out = +# in/2 must be integer-exact); the doubling families (upscale2/tile2) keep +# inputs in [3, 6] so their 2x outputs stay cheap under the full-budget ES fit +# (gather cost ~ out_cells x window area); everything else draws [4, 8]. +def rand_dim(name: String) -> Int: + if name == "subsample2": + return 4 + 2 * Int(random_float64(0.0, 3.0)) # {4, 6, 8} + if name == "upscale2" or name == "tile2": + return 3 + Int(random_float64(0.0, 4.0)) # [3, 6] + return 4 + Int(random_float64(0.0, 5.0)) # [4, 8] + + +# Ground-truth shape-changing transforms (what the engine must rediscover). +def apply_transform(name: String, g: ArcGrid) -> ArcGrid: + if name == "crop1": + var out = ArcGrid(g.rows - 2, g.cols - 2) + for r in range(out.rows): + for c in range(out.cols): + out.set(r, c, g.get(r + 1, c + 1)) + return out^ + elif name == "flip_h_crop1": + var out = ArcGrid(g.rows - 2, g.cols - 2) + for r in range(out.rows): + for c in range(out.cols): + out.set(r, c, g.get(r + 1, g.cols - 2 - c)) + return out^ + elif name == "subsample2": + # Every 2nd cell (even dims -> out = in/2 exactly). + var out = ArcGrid(g.rows // 2, g.cols // 2) + for r in range(out.rows): + for c in range(out.cols): + out.set(r, c, g.get(2 * r, 2 * c)) + return out^ + elif name == "upscale2": + # Blocky replication: each cell -> a 2x2 block. + var out = ArcGrid(g.rows * 2, g.cols * 2) + for r in range(out.rows): + for c in range(out.cols): + out.set(r, c, g.get(r // 2, c // 2)) + return out^ + else: # tile2: the grid replicated 2x2 (the modular sawtooth) + var out = ArcGrid(g.rows * 2, g.cols * 2) + for r in range(out.rows): + for c in range(out.cols): + out.set(r, c, g.get(r % g.rows, c % g.cols)) + return out^ + + +def make_demos(name: String) -> List[ArcTaskPair]: + var demos = List[ArcTaskPair]() + for _ in range(8): + var gin = rand_grid(rand_dim(name), rand_dim(name)) + var gout = apply_transform(name, gin) + demos.append(ArcTaskPair(gin^, gout^)) + return demos^ + + +# Held-out eval at a FRESH input size: the predicted output shape must match the +# truth (a wrong shape rule scores 0 for that trial — no OOB compare). With +# `plain` the content is read through the NON-toroidal gather instead (the wrap +# ablation: same fitted state, modular addressing removed). +def eval_held_out( + name: String, + state: UnsafePointer[Float32, MutAnyOrigin], + plain: Bool, +) raises -> Float32: + var match_sum = Float32(0.0) + var trials = 8 + for _ in range(trials): + var test_in = rand_grid(rand_dim(name), rand_dim(name)) + var truth = apply_transform(name, test_in) + var pr = ShapeGeomComposedMemory.out_rows(state, test_in) + var pc = ShapeGeomComposedMemory.out_cols(state, test_in) + if pr != truth.rows or pc != truth.cols: + continue # predicted shape wrong -> 0 for this trial + var pred = alloc[Float32](pr * pc) + if plain: + # The attention slots are state[0:7]; the plain gather ignores + # trel and reads a bounded (non-wrapped) source. + AttnGatherMemory.apply_shaped(state, test_in, pr, pc, pred) + else: + ShapeGeomComposedMemory.apply(state, test_in, pr, pc, pred) + match_sum += exact_match(pred, truth.data, pr * pc) + pred.free() + return match_sum / Float32(trials) + + +# Worst-case flat capacity over a demo list (the forward scratch size). +def demos_capacity(demos: List[ArcTaskPair]) -> Int: + var cap = 1 + for d in range(len(demos)): + if demos[d].input_grid.size() > cap: + cap = demos[d].input_grid.size() + if demos[d].output_grid.size() > cap: + cap = demos[d].output_grid.size() + return cap + + +# Cold per-task protocol: seed -> ONE fit_shape_geom call -> held-out eval. +# The RNG is re-seeded PER TASK (the arc_solve protocol) so a task's stochastic +# ES fit depends only on the task, not its position in the test. The caller +# owns `state` (seeded + fit here) so controls can re-read the fitted params. +def learn_and_eval( + name: String, + task_seed: Int, + state: UnsafePointer[Float32, MutAnyOrigin], +) raises -> Float32: + seed(task_seed) + var demos = make_demos(name) + var cap = demos_capacity(demos) + ShapeGeomComposedMemory.seed(state) + fit_shape_geom( + state, + demos, + cap, + FIT_N, + FIT_ALPHA0, + FIT_ALPHA1, + FIT_SIGMA0, + FIT_SIGMA1, + FIT_ITERS, + FIT_REG, + ) + return eval_held_out(name, state, False) + + +def main() raises: + seed(0) + + # ---- Ckpt A: the shape write recovers (k=1, b=-2) per axis on crop1. + var demos_a = make_demos("crop1") + var state_a = alloc[Float32](SHAPEGEOM_DIM) + ShapeGeomComposedMemory.seed(state_a) + ShapeGeomComposedMemory.write(state_a, demos_a) + var kr = state_a[SHAPEGEOM_SHAPE_OFF + 0] + var br = state_a[SHAPEGEOM_SHAPE_OFF + 1] + var kc = state_a[SHAPEGEOM_SHAPE_OFF + 2] + var bc = state_a[SHAPEGEOM_SHAPE_OFF + 3] + state_a.free() + var tol = Float32(0.05) + if ( + abs(kr - 1.0) > tol + or abs(br + 2.0) > tol + or abs(kc - 1.0) > tol + or abs(bc + 2.0) > tol + ): + raise Error( + "ERROR (Ckpt A): crop1 shape rule = (" + + String(kr) + + "*r+" + + String(br) + + ", " + + String(kc) + + "*c+" + + String(bc) + + "), expected (1*r-2, 1*c-2)." + ) + print("Ckpt A passed: shape write recovers crop1's (k=1, b=-2) per axis.") + + # ---- Ckpt B: the seam + geometry bar — each family cold at a fresh size. + # The wrap-ablation control below re-reads tile2's fitted params. + var names = List[String]() + names.append("crop1") + names.append("flip_h_crop1") + names.append("subsample2") + names.append("upscale2") + names.append("tile2") + + var wrap_ctl = Float32(-1.0) + var solved = 0 + for i in range(len(names)): + var state_b = alloc[Float32](SHAPEGEOM_DIM) + var held_out = learn_and_eval(names[i], i + 1, state_b) + print(" ", names[i], " held-out:", held_out) + if held_out >= 0.95: + solved += 1 + if names[i] == "tile2": + # Wrap ablation: the SAME fitted state through the plain gather. + wrap_ctl = eval_held_out("tile2", state_b, True) + state_b.free() + if solved != len(names): + raise Error( + "ERROR (Ckpt B): the shape memory did not solve the whole shape" + " family {crop, flip-crop, subsample, upscale, tile} to >= 0.95" + " held-out (" + + String(solved) + + "/" + + String(len(names)) + + " families)." + ) + print("Ckpt B passed: the whole shape family solved cold, held-out.") + + # ---- Control (wrap ablation): tile2's fitted params, but read through the + # NON-toroidal gather — the modular sawtooth `in[r mod n]` is provably + # outside any single affine (M, t), so removing the wrap must collapse it. + print(" control (plain gather) tile2 held-out:", wrap_ctl) + if wrap_ctl >= 0.5: + raise Error( + "ERROR (control): tile2 reached " + + String(wrap_ctl) + + " through the plain (non-toroidal) gather — the modular source" + " addressing should be load-bearing." + ) + print("Control passed: no wrap -> tiling collapses (as it must).") + + # ---- Control (shape ablation): the SAME content ES fit but with NO shape + # write (identity shape rule from the seed). Every demo's predicted output + # area then mismatches its true output area (fitness_shape penalizes them + # all -> no ES signal), and at eval the predicted shape is wrong -> ~0. + seed(7) + var demos_ctl = make_demos("crop1") + var cap_ctl = demos_capacity(demos_ctl) + var state_ctl = alloc[Float32](SHAPEGEOM_DIM) + ShapeGeomComposedMemory.seed(state_ctl) # identity shape rule, NOT written + # Mirror fit_shape_geom's discrete regime (hard frozen read) — the control + # ablates ONLY the shape write. + state_ctl[ATTN_BETA_OFF] = SHAPE_BETA_READ + var slow_ctl = alloc[Float32](SHAPEGEOM_DIM) + ShapeGeomComposedMemory.seed(slow_ctl) + fit_shape[ShapeGeomComposedMemory]( + state_ctl, + slow_ctl, + demos_ctl, + cap_ctl, + FIT_N, + FIT_ALPHA0, + FIT_ALPHA1, + FIT_SIGMA0, + FIT_SIGMA1, + FIT_ITERS, + FIT_REG, + ) + var held_out_ctl = eval_held_out("crop1", state_ctl, False) + slow_ctl.free() + state_ctl.free() + print(" control (no shape write) crop1 held-out:", held_out_ctl) + if held_out_ctl >= 0.5: + raise Error( + "ERROR (control): crop1 reached " + + String(held_out_ctl) + + " without the shape write — the ablation should fail, the shape" + " rule may not be load-bearing." + ) + print("Control passed: no shape rule -> wrong output size (as it must).") + + print( + "Shape-change test passed: the output-size seam works — output shape" + " inferred in-context, geometry fit on the output grid." + ) diff --git a/tests/test_shape_color.mojo b/tests/test_shape_color.mojo new file mode 100644 index 0000000..5940162 --- /dev/null +++ b/tests/test_shape_color.mojo @@ -0,0 +1,464 @@ +# suite-tier: full +from std.memory import alloc, UnsafePointer +from std.random import seed, random_float64 +from std.math import round +from std.collections import List, InlineArray + +# Run from the project root: `mojo run -I src tests/test_shape_color.mojo`. +from hope import ArcGrid, ArcTaskPair, COLOR_DIM +from memory_composed import ( + ShapeGeomColorComposedMemory, + SHAPEGEOMCOLOR_DIM, + SHAPEGEOMCOLOR_V_OFF, +) +from esper_evolution import ( + fit_shape_color, + FIT_N, + FIT_ALPHA0, + FIT_ALPHA1, + FIT_SIGMA0, + FIT_SIGMA1, + FIT_ITERS, + FIT_REG, +) +from arc_io import exact_match + +# ========================================================================== +# RUNG C — COLOUR ON TOP OF SHAPE. ShapeGeomColorComposedMemory composes a +# written colour table V on top of the proven shape+geometry gather: +# +# out = shape_geom_gather( V(in) ) +# +# Colour is cellwise (commutes with the copy gather — colour-then-gather), so V +# is written closed-form from the demos' FRACTION signatures (scale-invariant — +# the fix for count-conservation breaking under shape change) and the geometry +# is the unchanged two-frame multi-start fit on V-pre-mapped demos. +# +# THE COUNT-SIGNATURE WRITE'S HONEST PRECONDITION — colour-count CONTRAST. It +# matches colours by their across-demo count signatures, which identifies a +# permutation only when colours have DISTINGUISHABLE frequencies. Under EXACT +# same-shape conservation (GeomColor) even the tiny fluctuations of a uniform- +# random grid match exactly. Shape change is LOSSY (crop drops a border, +# subsample thins), so the signal must be REAL contrast, not fluctuation — which +# is exactly what real ARC grids have (a background + a few sparse colours, very +# different counts) and uniform-random grids lack. So these tasks draw grids +# from a per-task palette of DISTINCT per-colour frequencies (real-ARC-like); +# the uniform-random adversarial ceiling for crop is measured separately below. +# +# Ckpt A — the fraction colour write recovers the task's permutation on +# recolor_upscale2 (exact area ratio) AND recolor_crop1 (lossy border). +# Ckpt B — {recolor_crop1, recolor_subsample2, recolor_upscale2, recolor_tile2} +# each >= 0.95 held-out at a FRESH input size, per-task cold. +# Control (colour ablation) — the SAME fitted state with V forced to identity +# fails the recolored families: the colour module is load-bearing. +# Control (count-contrast ceiling) — the write on UNIFORM-random crop demos +# (no contrast) fails to recover V: the precondition is real, the +# contrast load-bearing (documented, not fought). +# Regression — a PURE-shape family (tile2, V = identity) through the new +# memory reproduces test_shape_change's bar (the strict superset). +# Few-demo — recolor_upscale2 at n=3 (the corpus median) clears the bar. +# +# Each task uses a RANDOM per-task palette (used colours + distinct frequencies) +# and a RANDOM permutation over the used colours (identity on the rest), held +# CONSTANT across the task's demos and re-derivable by re-seeding (so controls +# need no plumbing). Demos are drawn at VARYING sizes (shape rule identifiable); +# the held-out test is a fresh, unseen input size. +# ========================================================================== + +comptime N_USED = 5 # colours actually used per task (real ARC uses few) + + +# A per-task palette: `perm` (the colour permutation ground truth, identity on +# unused colours) and `w` (per-colour sampling weight, 0 for unused, DISTINCT +# for the used colours so their count signatures are separable). +struct Palette(Copyable, Movable): + var perm: InlineArray[Int, COLOR_DIM] + var w: InlineArray[Float32, COLOR_DIM] + + def __init__( + out self, + perm: InlineArray[Int, COLOR_DIM], + w: InlineArray[Float32, COLOR_DIM], + ): + self.perm = perm + self.w = w + + +# Draw a task palette off the current RNG: shuffle the colours, take the first +# N_USED as the used set, permute them cyclically (a derangement — every used +# colour actually changes), and give them distinct descending weights. +def rand_palette() -> Palette: + var idx = InlineArray[Int, COLOR_DIM](fill=0) + for i in range(COLOR_DIM): + idx[i] = i + for i in range(COLOR_DIM - 1, 0, -1): + var j = Int(random_float64(0.0, Float64(i + 1))) + var tmp = idx[i] + idx[i] = idx[j] + idx[j] = tmp + + var perm = InlineArray[Int, COLOR_DIM](fill=0) + for i in range(COLOR_DIM): + perm[i] = i # identity on unused colours + var w = InlineArray[Float32, COLOR_DIM](fill=0.0) + for i in range(N_USED): + perm[idx[i]] = idx[ + (i + 1) % N_USED + ] # cyclic permutation of the used set + w[idx[i]] = Float32(N_USED - i + 1) # distinct weights: 6, 5, 4, 3, 2 + return Palette(perm, w) + + +# A grid whose cells are drawn from the palette's weighted colour distribution +# (categorical sampling) — distinct per-colour frequencies, real-ARC-like. +def rand_grid( + rows: Int, cols: Int, w: InlineArray[Float32, COLOR_DIM] +) -> ArcGrid: + var total = Float32(0.0) + for c in range(COLOR_DIM): + total += w[c] + var g = ArcGrid(rows, cols) + for k in range(rows * cols): + var u = Float32(random_float64(0.0, Float64(total))) + var acc = Float32(0.0) + var col = 0 + for c in range(COLOR_DIM): + acc += w[c] + if u < acc: + col = c + break + g.data[k] = Float32(col) + return g^ + + +# A uniform-random grid (the adversarial no-contrast case for the ceiling +# control): every colour equiprobable, so count signatures barely separate. +def rand_grid_uniform(rows: Int, cols: Int) -> ArcGrid: + var g = ArcGrid(rows, cols) + for k in range(rows * cols): + g.data[k] = Float32(Int(random_float64(0.0, Float64(COLOR_DIM)))) + return g^ + + +def rand_dim(name: String) -> Int: + if name == "recolor_subsample2": + # Even dims {8, 10, 12}: subsample HALVES each axis, so the output must + # be large enough for a reliable colour signature (a 2x2 output starves + # the low-frequency colours — real-ARC subsample tasks aren't 4x4). + return 8 + 2 * Int(random_float64(0.0, 3.0)) # {8, 10, 12} + if name == "recolor_upscale2" or name == "recolor_tile2" or name == "tile2": + return 3 + Int( + random_float64(0.0, 4.0) + ) # [3, 6] (doubling stays cheap) + return 4 + Int(random_float64(0.0, 5.0)) # [4, 8] (crop1, recolor_crop1) + + +# A pure-shape stress family draws UNIFORM-random grids (no colour contrast) so +# the count-signature write has NO recolor signal — the adversarial case where a +# naive write scrambles V. The global recolor gate must keep V = identity here +# (strict superset). "tile2" keeps palette grids (exact-conservation regression). +def is_uniform_pure(name: String) -> Bool: + return name == "crop1" + + +# Ground-truth colour-on-shape transform: recolor each cell through `perm`, then +# apply the shape change. (Recolor commutes with the copy gather.) The pure +# "tile2" regression control skips the recolor so its written V is identity. +def apply_transform( + name: String, g: ArcGrid, perm: InlineArray[Int, COLOR_DIM] +) -> ArcGrid: + var rc = ArcGrid(g.rows, g.cols) + for k in range(g.rows * g.cols): + if name == "tile2" or name == "crop1": + rc.data[k] = g.data[k] # pure shape: no recolor + else: + rc.data[k] = Float32(perm[Int(round(g.data[k]))]) + + if name == "recolor_crop1" or name == "crop1": + var out = ArcGrid(rc.rows - 2, rc.cols - 2) + for r in range(out.rows): + for c in range(out.cols): + out.set(r, c, rc.get(r + 1, c + 1)) + return out^ + elif name == "recolor_subsample2": + var out = ArcGrid(rc.rows // 2, rc.cols // 2) + for r in range(out.rows): + for c in range(out.cols): + out.set(r, c, rc.get(2 * r, 2 * c)) + return out^ + elif name == "recolor_upscale2": + var out = ArcGrid(rc.rows * 2, rc.cols * 2) + for r in range(out.rows): + for c in range(out.cols): + out.set(r, c, rc.get(r // 2, c // 2)) + return out^ + else: # recolor_tile2 or tile2 (both tile; tile2's rc is un-recolored) + var out = ArcGrid(rc.rows * 2, rc.cols * 2) + for r in range(out.rows): + for c in range(out.cols): + out.set(r, c, rc.get(r % rc.rows, c % rc.cols)) + return out^ + + +def make_demos(name: String, n: Int, pal: Palette) -> List[ArcTaskPair]: + var demos = List[ArcTaskPair]() + for _ in range(n): + var r = rand_dim(name) + var c = rand_dim(name) + var gin = rand_grid_uniform(r, c) if is_uniform_pure( + name + ) else rand_grid(r, c, pal.w) + var gout = apply_transform(name, gin, pal.perm) + demos.append(ArcTaskPair(gin^, gout^)) + return demos^ + + +# Held-out eval at a FRESH input size (same palette). A wrong predicted shape +# scores 0 for that trial. `ablate_v` forces V = identity onto a COPY of the +# fitted state (the colour-ablation control). +def eval_held_out( + name: String, + state: UnsafePointer[Float32, MutAnyOrigin], + pal: Palette, + ablate_v: Bool, +) raises -> Float32: + var st = alloc[Float32](SHAPEGEOMCOLOR_DIM) + for i in range(SHAPEGEOMCOLOR_DIM): + st[i] = state[i] + if ablate_v: + for s in range(COLOR_DIM): + st[SHAPEGEOMCOLOR_V_OFF + s] = Float32(s) + + var match_sum = Float32(0.0) + var trials = 8 + for _ in range(trials): + var tr = rand_dim(name) + var tc = rand_dim(name) + var test_in = rand_grid_uniform(tr, tc) if is_uniform_pure( + name + ) else rand_grid(tr, tc, pal.w) + var truth = apply_transform(name, test_in, pal.perm) + var pr = ShapeGeomColorComposedMemory.out_rows(st, test_in) + var pc = ShapeGeomColorComposedMemory.out_cols(st, test_in) + if pr != truth.rows or pc != truth.cols: + continue # predicted shape wrong -> 0 for this trial + var pred = alloc[Float32](pr * pc) + ShapeGeomColorComposedMemory.apply(st, test_in, pr, pc, pred) + match_sum += exact_match(pred, truth.data, pr * pc) + pred.free() + st.free() + return match_sum / Float32(trials) + + +def demos_capacity(demos: List[ArcTaskPair]) -> Int: + var cap = 1 + for d in range(len(demos)): + if demos[d].input_grid.size() > cap: + cap = demos[d].input_grid.size() + if demos[d].output_grid.size() > cap: + cap = demos[d].output_grid.size() + return cap + + +# Cold per-task protocol: seed -> ONE fit_shape_color -> held-out eval. RNG +# re-seeded PER TASK (the arc_solve protocol) so the fit depends only on the +# task; the palette is re-derivable by re-seeding task_seed. Caller owns `state`. +def learn_and_eval( + name: String, + task_seed: Int, + n_demos: Int, + state: UnsafePointer[Float32, MutAnyOrigin], +) raises -> Float32: + seed(task_seed) + var pal = rand_palette() + var demos = make_demos(name, n_demos, pal) + var cap = demos_capacity(demos) + ShapeGeomColorComposedMemory.seed(state) + fit_shape_color( + state, + demos, + cap, + FIT_N, + FIT_ALPHA0, + FIT_ALPHA1, + FIT_SIGMA0, + FIT_SIGMA1, + FIT_ITERS, + FIT_REG, + ) + return eval_held_out(name, state, pal, False) + + +# Count how many of the palette's colours the written V got wrong (write-only — +# no ES fit, so cheap; the geometry-invariant colour write is all that runs). +def write_miss_count(name: String, task_seed: Int, uniform: Bool) raises -> Int: + seed(task_seed) + var pal = rand_palette() + var demos = List[ArcTaskPair]() + for _ in range(8): + var r = rand_dim(name) + var c = rand_dim(name) + var gin = rand_grid_uniform(r, c) if uniform else rand_grid(r, c, pal.w) + var gout = apply_transform(name, gin, pal.perm) + demos.append(ArcTaskPair(gin^, gout^)) + var st = alloc[Float32](SHAPEGEOMCOLOR_DIM) + ShapeGeomColorComposedMemory.seed(st) + ShapeGeomColorComposedMemory.write(st, demos) + var wrong = 0 + # A uniform grid uses all colours; a palette grid only its used set. Compare + # only colours the demos actually contain (unseen default to identity). + for c in range(COLOR_DIM): + if uniform or pal.w[c] > 0.0: + if Int(round(st[SHAPEGEOMCOLOR_V_OFF + c])) != pal.perm[c]: + wrong += 1 + st.free() + return wrong + + +def main() raises: + seed(0) + + # ---- Ckpt A: the fraction colour write recovers the task permutation on + # both the exact-ratio (upscale) and lossy (crop) families, given contrast. + var a_names = List[String]() + a_names.append("recolor_upscale2") + a_names.append("recolor_crop1") + a_names.append("recolor_subsample2") + for ai in range(len(a_names)): + var wrong = write_miss_count(a_names[ai], 100 + ai, False) + if wrong != 0: + raise Error( + "ERROR (Ckpt A): " + + a_names[ai] + + " colour write missed " + + String(wrong) + + " used colours of the task permutation." + ) + print("Ckpt A passed: the fraction colour write recovers the permutation.") + + # ---- Control (count-contrast ceiling): the SAME write on UNIFORM-random + # crop demos (no contrast) should FAIL to recover V — documenting that the + # count-signature precondition is real, not scaffolding. + var uni_wrong = write_miss_count("recolor_crop1", 100, True) + print( + " control (uniform-random crop, no contrast) colours missed:", + uni_wrong, + "/ 10", + ) + if uni_wrong == 0: + raise Error( + "ERROR (control): the colour write recovered V from UNIFORM-random" + " crop demos — the count-contrast precondition should be load-" + "bearing under lossy shape change." + ) + print( + "Control passed: no contrast -> crop colour write fails (as it must)." + ) + + # ---- Ckpt B: each colour-on-shape family cold at a fresh size. + var names = List[String]() + names.append("recolor_crop1") + names.append("recolor_subsample2") + names.append("recolor_upscale2") + names.append("recolor_tile2") + + var ablate_ctl = Float32(-1.0) + var solved = 0 + for i in range(len(names)): + var state_b = alloc[Float32](SHAPEGEOMCOLOR_DIM) + var held_out = learn_and_eval(names[i], i + 1, 8, state_b) + print(" ", names[i], " held-out:", held_out) + if held_out >= 0.95: + solved += 1 + if names[i] == "recolor_upscale2": + # Colour ablation: the SAME fitted state, V forced to identity. + # Re-derive the task palette by re-seeding (deterministic). + seed(i + 1) + var pal_b = rand_palette() + ablate_ctl = eval_held_out(names[i], state_b, pal_b, True) + state_b.free() + if solved != len(names): + raise Error( + "ERROR (Ckpt B): the colour-on-shape family was not solved to" + " >= 0.95 held-out (" + + String(solved) + + "/" + + String(len(names)) + + " families)." + ) + print("Ckpt B passed: the whole colour-on-shape family solved cold.") + + # ---- Control (colour ablation): fitted recolor_upscale2 with V = identity. + print(" control (V = identity) recolor_upscale2 held-out:", ablate_ctl) + if ablate_ctl >= 0.5: + raise Error( + "ERROR (control): recolor_upscale2 reached " + + String(ablate_ctl) + + " with V = identity — the colour module should be load-bearing." + ) + print("Control passed: no colour table -> recolor collapses (as it must).") + + # ---- Regression: a PURE-shape family (tile2) through the new memory. With + # no recolor the written V is identity, so this must reproduce + # test_shape_change's tile2 bar (the strict-superset guard). + var state_r = alloc[Float32](SHAPEGEOMCOLOR_DIM) + var held_r = learn_and_eval("tile2", 5, 8, state_r) + state_r.free() + print(" regression (pure tile2 through colour memory) held-out:", held_r) + if held_r < 0.95: + raise Error( + "ERROR (regression): pure tile2 through the colour memory scored " + + String(held_r) + + " (< 0.95) — the identity-V path should match the shape memory." + ) + print("Regression passed: V = identity is the pure shape path.") + + # ---- Regression (the strict-superset guard): a PURE-shape CROP on + # UNIFORM-random grids (no colour contrast) — the adversarial case where a + # naive count write scrambles V (measured 1.0 -> 0.17). The GLOBAL recolor + # gate must keep V = identity, so held-out stays at the pure-shape bar. + var state_c = alloc[Float32](SHAPEGEOMCOLOR_DIM) + var held_c = learn_and_eval("crop1", 11, 8, state_c) + var v_wrong = 0 + for c in range(COLOR_DIM): + if Int(round(state_c[SHAPEGEOMCOLOR_V_OFF + c])) != c: + v_wrong += 1 + state_c.free() + print( + " regression (pure UNIFORM crop1) held-out:", + held_c, + " V non-identity entries:", + v_wrong, + ) + if held_c < 0.95 or v_wrong != 0: + raise Error( + "ERROR (strict superset): pure uniform crop1 scored " + + String(held_c) + + " with " + + String(v_wrong) + + " non-identity V entries — the recolor gate must keep V =" + " identity" + " on a no-contrast pure-shape task (else the colour path regresses" + " the pure-shape path)." + ) + print("Regression passed: no-contrast pure crop keeps V = identity.") + + # ---- Few-demo (corpus median n=3): the fraction-write must still recover V + # and clear the bar at 3 demos. + var state_f = alloc[Float32](SHAPEGEOMCOLOR_DIM) + var held_f = learn_and_eval("recolor_upscale2", 9, 3, state_f) + state_f.free() + print(" few-demo (recolor_upscale2, n=3) held-out:", held_f) + if held_f < 0.95: + raise Error( + "ERROR (few-demo): recolor_upscale2 at n=3 scored " + + String(held_f) + + " (< 0.95) — the fraction colour write should hold at the corpus" + " median 3 demos." + ) + print("Few-demo passed: the fraction colour write holds at n=3.") + + print( + "Colour-on-shape test passed (Rung C): the shape path now composes a" + " written colour table — held-out, cold, at fresh sizes." + ) diff --git a/tools/synth_tasks.py b/tools/synth_tasks.py index dc441a4..17a7f41 100644 --- a/tools/synth_tasks.py +++ b/tools/synth_tasks.py @@ -65,6 +65,53 @@ def _shift(grid): return [[row[-1]] + row[:-1] for row in grid] +# --------------------------------------------------------------------------- +# SHAPE-CHANGING transforms (Vision A / Next #1). Unlike everything above these +# return a grid whose dims DIFFER from the input — the output shape is a rule to +# be inferred in-context (never a hand-coded size heuristic). They are the +# ground truth the engine's ShapeMemory must rediscover: a closed-form shape +# rule (out = k*in + b per axis) composed with the AttnGather content gather. +# --------------------------------------------------------------------------- +def _crop1(grid): + # Drop the 1-cell border: (r-2, c-2). Centred crop -> identity content. + return [row[1:-1] for row in grid[1:-1]] + + +def _flip_h_crop1(grid): + # Crop the 1-cell border, then reverse columns (a flip within the resize). + return [list(reversed(row[1:-1])) for row in grid[1:-1]] + + +def _subsample2(grid): + # Take every 2nd cell on each axis: (ceil(r/2), ceil(c/2)). Exactly r/2, c/2 + # for even dims (the shape rule k=1/2, b=0 is then integer-exact). + return [ + [grid[i][j] for j in range(0, len(grid[0]), 2)] + for i in range(0, len(grid), 2) + ] + + +def _upscale2(grid): + # Blocky replication: each cell becomes a 2x2 block -> (2r, 2c). The shape + # rule is k=2, b=0; the content is out[r][c] = in[r//2][c//2] (a floor + # gather, which the affine attention gather expresses exactly at sharp + # temperature: M = I/2, t = 0). + return [ + [grid[i // 2][j // 2] for j in range(2 * len(grid[0]))] + for i in range(2 * len(grid)) + ] + + +def _tile2(grid): + # Tile the grid 2x2 -> (2r, 2c). Same shape rule as upscale2 (k=2, b=0) but + # the content out[r][c] = in[r % rows][c % cols] is a sawtooth — genuinely + # NON-affine, the family that forces modular (wrapped) source addressing. + return [ + [grid[i % len(grid)][j % len(grid[0])] for j in range(2 * len(grid[0]))] + for i in range(2 * len(grid)) + ] + + TRANSFORMS = { "identity": _identity, "flip_h": _flip_h, @@ -74,6 +121,42 @@ def _shift(grid): "shift": _shift, } +# Shape-changing families kept in a SEPARATE table: they need their own group +# generator (input size must VARY across a task's demos so the shape rule is +# identifiable, not memorized). `subsample2` requires even input dims. +# COLOUR-ON-SHAPE families (Vision A / Next #1, Rung C): a shape change composed +# with a cellwise recolor. Colour commutes with the copy gather, so applying +# `_recolor` (the cyclic +1 palette shift) before the shape transform is the +# same as after — the ground truth the ShapeGeomColorComposedMemory must +# rediscover as (shape rule, written colour table V, geometry). +def _recolor_crop1(grid): + return _crop1(_recolor(grid)) + + +def _recolor_subsample2(grid): + return _subsample2(_recolor(grid)) + + +def _recolor_upscale2(grid): + return _upscale2(_recolor(grid)) + + +def _recolor_tile2(grid): + return _tile2(_recolor(grid)) + + +SHAPE_TRANSFORMS = { + "crop1": _crop1, + "flip_h_crop1": _flip_h_crop1, + "subsample2": _subsample2, + "upscale2": _upscale2, + "tile2": _tile2, + "recolor_crop1": _recolor_crop1, + "recolor_subsample2": _recolor_subsample2, + "recolor_upscale2": _recolor_upscale2, + "recolor_tile2": _recolor_tile2, +} + def _random_grid(rows, cols, rng): return [[rng.randrange(NUM_COLORS) for _ in range(cols)] for _ in range(rows)] @@ -148,6 +231,54 @@ def generate_task_groups(transform, out_dir, num_tasks, n_train, rows, cols, see return paths +def _rand_shape_size(rng, even): + """A random grid size in [4, 8] per axis; even-only when `even` (subsample).""" + if even: + r = 4 + 2 * rng.randrange(3) # {4, 6, 8} + c = 4 + 2 * rng.randrange(3) + else: + r = 4 + rng.randrange(5) # [4, 8] + c = 4 + rng.randrange(5) + return r, c + + +def generate_shape_task_groups(transform, out_dir, num_tasks, n_train, seed): + """Emit `num_tasks` SHAPE-CHANGING task bundles (`.task`) for `transform`. + + Unlike `generate_task_groups`, each demo (and the held-out test) is drawn at + a RANDOM input size, so the per-axis shape rule out = k*in + b is genuinely + IDENTIFIABLE from the demos (>= 2 distinct sizes) rather than memorized, and + the unseen-size test pair is an uncheatable generalization probe. Returns the + bundle paths written. + """ + if transform not in SHAPE_TRANSFORMS: + raise ValueError( + "Unknown shape transform %r; choose from %s" + % (transform, ", ".join(sorted(SHAPE_TRANSFORMS))) + ) + os.makedirs(out_dir, exist_ok=True) + fn = SHAPE_TRANSFORMS[transform] + even = transform in ("subsample2", "recolor_subsample2") + rng = _random.Random(seed) + + paths = [] + for t in range(num_tasks): + train = [] + for _ in range(n_train): + r, c = _rand_shape_size(rng, even) + grid_in = _random_grid(r, c, rng) + train.append((grid_in, fn(grid_in))) + r, c = _rand_shape_size(rng, even) + test_in = _random_grid(r, c, rng) + test = [(test_in, fn(test_in))] + + path = os.path.join(out_dir, "%s_%d.task" % (transform, t)) + _save_task(train, test, path) + paths.append(path) + + return paths + + def _parse_args(): p = argparse.ArgumentParser(description=__doc__) p.add_argument(