esper_viz.py

import marimo

__generated_with = "0.23.11"
app = marimo.App(width="medium")


@app.cell
def _(mo):
    mo.md(r"""
    # Esper — learning a grid transform from a handful of demos

    Esper is a bare-metal, pure-[Mojo](https://www.modular.com/mojo)
    reasoning engine. Given a few **demonstration pairs** (`input → output`),
    it fits a small *operator* to them in-context — using a derivative-free
    **Evolution Strategy**, no backprop, no neural-net framework — and then
    applies that operator to an **unseen test input**.

    The number that matters is **held-out exact match**: fit on the train
    pairs, score the *test* pair the model never saw. It is uncheatable by
    memorization. Below, each task shows the test **input**, Esper's
    **prediction**, and the **truth** — pick a task to see it generalize.

    *These panels are rendered from `arc_solve --trace` output; nothing here
    is hand-drawn.*
    """)
    return


@app.cell
def _():
    import anywidget
    import marimo as mo
    import traitlets

    return anywidget, mo, traitlets


@app.cell
def _():
    # The ARC colour palette (values 0-9). Defined HERE as the single source of
    # truth for colour across all of Esper's visual artifacts.
    PALETTE = [
        "#101014",  # 0 background (near-black)
        "#1E93FF",  # 1 blue
        "#F93C31",  # 2 red
        "#4FCC30",  # 3 green
        "#FFDC00",  # 4 yellow
        "#999999",  # 5 grey
        "#E53AA3",  # 6 magenta
        "#FF851B",  # 7 orange
        "#87D8F1",  # 8 sky
        "#921231",  # 9 maroon
    ]
    return (PALETTE,)


@app.cell
def _(anywidget, traitlets):
    class ArcGridWidget(anywidget.AnyWidget):
        """An ARC grid: a 2-D array of colour indices, rendered as cells."""

        _esm = """
        function render({ model, el }) {
          function draw() {
            const grid = model.get("grid") || [];
            const palette = model.get("palette") || [];
            const cell = model.get("cell_px") || 22;
            const rows = grid.length;
            const cols = rows ? grid[0].length : 0;
            el.innerHTML = "";
            const wrap = document.createElement("div");
            wrap.className = "arc-grid";
            wrap.style.gridTemplateColumns = `repeat(${cols}, ${cell}px)`;
            wrap.style.gridTemplateRows = `repeat(${rows}, ${cell}px)`;
            for (let r = 0; r < rows; r++) {
              for (let c = 0; c < cols; c++) {
                const v = grid[r][c];
                const d = document.createElement("div");
                d.className = "arc-cell";
                d.style.background = palette[v] ?? "#888";
                wrap.appendChild(d);
              }
            }
            el.appendChild(wrap);
          }
          draw();
          model.on("change:grid", draw);
          model.on("change:palette", draw);
        }
        export default { render };
        """
        _css = """
        .arc-grid {
          display: grid;
          gap: 1px;
          background: #cfcfcf;
          padding: 2px;
          border-radius: 5px;
          width: max-content;
        }
        .arc-cell { width: 100%; height: 100%; border-radius: 1px; }
        @media (prefers-color-scheme: dark) {
          .arc-grid { background: #3a3a3a; }
        }
        """

        grid = traitlets.List([]).tag(sync=True)
        palette = traitlets.List([]).tag(sync=True)
        cell_px = traitlets.Int(22).tag(sync=True)

    return (ArcGridWidget,)


@app.cell
def _():
    # A self-contained sample of `arc_solve --trace` output (real held-out
    # solves; the WASM export needs no external file). Regenerate with:
    #   ./esper run src/arc_solve.mojo --report --trace out.jsonl <tasks>
    TRACE = [
        {
            "name": "flip_h · #0",
            "mem": "same",
            "held_out": 1.0,
            "train_fit": 1.0,
            "input": [[1, 8, 6, 8], [4, 8, 3, 3], [9, 6, 9, 4], [7, 7, 5, 1]],
            "pred": [[8, 6, 8, 1], [3, 3, 8, 4], [4, 9, 6, 9], [1, 5, 7, 7]],
            "true": [[8, 6, 8, 1], [3, 3, 8, 4], [4, 9, 6, 9], [1, 5, 7, 7]],
        },
        {
            "name": "flip_h · #1",
            "mem": "same",
            "held_out": 1.0,
            "train_fit": 1.0,
            "input": [[0, 8, 4, 2], [3, 7, 5, 9], [4, 5, 9, 9], [2, 4, 6, 6]],
            "pred": [[2, 4, 8, 0], [9, 5, 7, 3], [9, 9, 5, 4], [6, 6, 4, 2]],
            "true": [[2, 4, 8, 0], [9, 5, 7, 3], [9, 9, 5, 4], [6, 6, 4, 2]],
        },
        {
            "name": "crop1 (shape-change) · #0",
            "mem": "shape",
            "held_out": 1.0,
            "train_fit": 1.0,
            "input": [
                [3, 1, 3, 0, 8],
                [7, 7, 4, 8, 6],
                [3, 3, 6, 6, 8],
                [0, 9, 9, 0, 6],
            ],
            "pred": [[7, 4, 8], [3, 6, 6]],
            "true": [[7, 4, 8], [3, 6, 6]],
        },
    ]
    return (TRACE,)


@app.cell
def _(ArcGridWidget, PALETTE, mo):
    def panel(title, grid, cell_px=26):
        w = ArcGridWidget(grid=grid, palette=PALETTE, cell_px=cell_px)
        return mo.vstack(
            [mo.md(f"**{title}**"), mo.ui.anywidget(w)], align="center"
        )

    def diff_cells(a, b):
        # Count mismatched cells (shape mismatch => all cells count).
        if len(a) != len(b) or (a and len(a[0]) != len(b[0])):
            return max(len(a) * (len(a[0]) if a else 0), 1)
        return sum(
            1
            for r in range(len(a))
            for c in range(len(a[0]))
            if a[r][c] != b[r][c]
        )

    return diff_cells, panel


@app.cell
def _(TRACE, mo):
    task_picker = mo.ui.dropdown(
        options={rec["name"]: i for i, rec in enumerate(TRACE)},
        value=TRACE[0]["name"],
        label="Task",
    )
    task_picker
    return (task_picker,)


@app.cell
def _(TRACE, diff_cells, mo, panel, task_picker):
    rec = TRACE[task_picker.value]
    _mismatch = diff_cells(rec["pred"], rec["true"])
    _verdict = (
        "✅ exact match — generalized to the unseen input"
        if _mismatch == 0
        else f"❌ {_mismatch} mismatched cell(s)"
    )
    mo.vstack(
        [
            mo.hstack(
                [
                    panel("test input", rec["input"]),
                    panel("Esper prediction", rec["pred"]),
                    panel("truth", rec["true"]),
                ],
                justify="center",
                gap=2,
            ),
            mo.md(
                f"held-out exact match: **{rec['held_out']:.2f}** &nbsp;·&nbsp; "
                f"train fit: **{rec['train_fit']:.2f}** &nbsp;·&nbsp; "
                f"memory: `{rec['mem']}` &nbsp;·&nbsp; {_verdict}"
            ),
        ]
    )
    return


if __name__ == "__main__":
    app.run()