atom-ai / test /test_decode_fix_bptt.py

ATOM (atom-ai): code snapshot (100fae0 + batched-train), README with measured numbers, curated checkpoints (part 4)

b9301f1 verified about 1 month ago

21.8 kB

"""Unit tests: byte-tick decode label fix + opt-in truncated BPTT through the field
(docs/DECODE_FIX_AND_BPTT.md).
Decode fix (automatic in byte-tick, escape flag ``legacy_gen_labels`` / --legacy-gen-labels):
* byte-tick generated commits are bit-identical (label + features + atomizer state) to
  ``encode_bytes`` of the same byte stream, in both pack modes;
* the teacher-forced and generation paths give identical logits (max |dlogit| == 0);
* ``legacy_gen_labels=True`` reproduces the pre-fix commit (commit f8a8918) exactly:
  labels, features, and a sampled generate_packets() byte string;
* non-byte-tick (max_span > 1) commits are unchanged whatever the flag.
Field BPTT (--field-bptt K, --field-leak-learnable):
* off (default): multi-tick train+eval trace bit-identical to commit f8a8918
  (no leak, scalar leak 0.3, leak spectrum), trainer checkpoint bit-identical too;
* on: forward values (field, logits, loss) bit-identical to off at equal weights;
  eval/no_grad ticks never keep a graph;
* on: d loss_t / d atom_{t-L} != 0 for L = 1..K-1, and == 0 across a window cut;
  off: == 0 for every L >= 1 (the ROOT_CAUSE measurement);
* learnable leak: init == scalar leak (bit-identical trace), receives gradient under BPTT;
* checkpoint roundtrip; old checkpoints load as off; trainer CLI default off, on,
  and CLI wins after --resume in both directions.
"""
from __future__ import annotations

import importlib
import shutil
import subprocess
import sys
import tempfile
import unittest
from pathlib import Path

import torch

from src.atom_native import AtomNativeModel
from src.io.atomizer import Atomizer

ROOT = Path(__file__).resolve().parents[1]
LEGACY_COMMIT = "f8a8918"
TEXT = "Bonjour, comment ça va ? Il était une fois un petit atome. Qui es-tu ?\nOui, d'accord... " * 3

_LEG: dict = {}


def _legacy():
    """(AtomNativeModel, Atomizer) classes of the pre-change commit, or None."""
    if "mods" in _LEG:
        return _LEG["mods"]
    try:
        tmp = Path(tempfile.mkdtemp(prefix="legacy_dfb_"))
        pkg = tmp / "legacy_src_dfb"
        pkg.mkdir()
        files = subprocess.run(["git", "-C", str(ROOT), "ls-tree", "-r", "--name-only", LEGACY_COMMIT, "src"],
                               capture_output=True, text=True, check=True).stdout.split()
        for f in files:
            if not f.endswith(".py"):
                continue
            code = subprocess.run(["git", "-C", str(ROOT), "show", f"{LEGACY_COMMIT}:{f}"],
                                  capture_output=True, text=True, check=True).stdout
            dst = pkg / Path(f).relative_to("src")
            dst.parent.mkdir(parents=True, exist_ok=True)
            dst.write_text(code, encoding="utf-8")
        for d in [pkg, *[p for p in pkg.rglob("*") if p.is_dir()]]:
            (d / "__init__.py").touch(exist_ok=True)
        sys.path.insert(0, str(tmp))
        try:
            an = importlib.import_module("legacy_src_dfb.atom_native")
            at = importlib.import_module("legacy_src_dfb.io.atomizer")
        finally:
            sys.path.remove(str(tmp))
        _LEG["mods"] = (an.AtomNativeModel, at.Atomizer)
    except Exception:  # pragma: no cover
        _LEG["mods"] = None
    return _LEG["mods"]


def _build(cls=AtomNativeModel, atom_cls=Atomizer, **kw):
    torch.manual_seed(0)
    return cls(
        d_model=8, n_modes=8, n_atoms_max=32, max_payload_bytes=1,
        atomizer=atom_cls(max_span_bytes=1), field_max_rms=3.0,
        field_obligatory_hard=True, enable_merge=False, last_atom_readout=True,
        field_loss_weight=0.05, **kw,
    )


def _run(model, n_train=40, n_eval=20):
    opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=1e-3)
    packets = model.atomizer.encode(TEXT, reset=True)
    model.reset_state(reset_atomizer=False)
    model.train()
    trace = []
    torch.manual_seed(5)
    for cur, tgt in list(zip(packets[:-1], packets[1:]))[:n_train]:
        opt.zero_grad(set_to_none=True)
        loss, _ = model.transition_loss(cur, tgt)
        loss.backward()
        opt.step()
        model.stabilize_dynamics_parameters()
        trace.append(loss.detach().clone())
        trace.append(model.core.state.alpha.detach().clone())
    model.eval()
    with torch.no_grad():
        for cur, tgt in list(zip(packets[:-1], packets[1:]))[n_train:n_train + n_eval]:
            loss, _ = model.transition_loss(cur, tgt)
            trace.append(loss.detach().clone())
            trace.append(model.core.state.alpha.detach().clone())
    trace.extend(p.detach().clone() for p in model.parameters())
    return trace


def _eq(tc, a, b):
    tc.assertEqual(len(a), len(b))
    for x, y in zip(a, b):
        tc.assertTrue(torch.equal(x, y))


class DecodeLabelFixTests(unittest.TestCase):
    def test_byte_tick_commit_equals_training_encode(self) -> None:
        raw = TEXT.encode("utf-8")
        for mode in ("linguistic", "eager"):
            tr = Atomizer(max_span_bytes=1, pack_mode=mode).encode_bytes(raw, reset=True, flush=False)
            gen_at = Atomizer(max_span_bytes=1, pack_mode=mode)
            gen = [gen_at.packet_from_payload(bytes([b])) for b in raw]
            self.assertEqual(len(tr), len(gen))
            for a, b in zip(tr, gen):
                self.assertEqual((a.payload, a.start, a.boundary, a.level), (b.payload, b.start, b.boundary, b.level))
                self.assertTrue(torch.equal(a.features, b.features))
            self.assertEqual({p.boundary for p in gen} - {"max_span", "newline", "punctuation", "whitespace"}, set())
            if mode == "linguistic":
                self.assertEqual({p.boundary for p in gen}, {"max_span", "newline"})

    def test_legacy_flag_reproduces_old_labels_exactly(self) -> None:
        leg = _legacy()
        if leg is None:
            self.skipTest("legacy commit not available")
        _, OldAtomizer = leg
        raw = TEXT.encode("utf-8")
        for mode in ("linguistic", "eager"):
            new = Atomizer(max_span_bytes=1, pack_mode=mode); new.legacy_gen_labels = True
            old = OldAtomizer(max_span_bytes=1, pack_mode=mode)
            for b in raw:
                a, o = new.packet_from_payload(bytes([b])), old.packet_from_payload(bytes([b]))
                self.assertEqual((a.boundary, a.start, a.level), (o.boundary, o.start, o.level))
                self.assertTrue(torch.equal(a.features, o.features))
        # the fix does change labels (else the test above would prove nothing)
        fixed = Atomizer(max_span_bytes=1)
        self.assertEqual(fixed.packet_from_payload(b" ").boundary, "max_span")
        legacy = Atomizer(max_span_bytes=1); legacy.legacy_gen_labels = True
        self.assertEqual(legacy.packet_from_payload(b" ").boundary, "whitespace")

    def test_non_byte_tick_unchanged(self) -> None:
        leg = _legacy()
        if leg is None:
            self.skipTest("legacy commit not available")
        _, OldAtomizer = leg
        for flag in (False, True):
            new = Atomizer(max_span_bytes=8); new.legacy_gen_labels = flag
            old = OldAtomizer(max_span_bytes=8)
            for pl in (b"abc ", b"Oui,", b"x", b" ", b"\n", b"mot"):
                a, o = new.packet_from_payload(pl), old.packet_from_payload(pl)
                self.assertEqual(a.boundary, o.boundary)
                self.assertTrue(torch.equal(a.features, o.features))

    def test_teacher_forced_equals_generation_path(self) -> None:
        m = _build(field_leak=0.3)
        m.eval()
        text = "Utilisateur: Qui es-tu ?\nAssistant: Je suis Atom, un petit modèle, d'accord ?\n"
        raw = text.encode()
        with torch.no_grad():
            m.reset_state(reset_atomizer=True)
            pk = m.atomizer.encode_bytes(raw, reset=True)
            tf = [m.forward_packet(p)["surface"]["byte_logits"].clone() for p in pk[:-1]]
            m.reset_state(reset_atomizer=True)
            n0 = 30
            pp = m.atomizer.encode_bytes(raw[:n0], reset=True)
            gen = [m.forward_packet(p)["surface"]["byte_logits"].clone() for p in pp[:-1]]
            cur = pp[-1]
            for b in raw[n0:]:
                gen.append(m.forward_packet(cur)["surface"]["byte_logits"].clone())
                cur = m.atomizer.packet_from_payload(bytes([b]))
        self.assertEqual(len(tf), len(gen))
        self.assertEqual(max(float((a - b).abs().max()) for a, b in zip(tf, gen)), 0.0)
        # legacy labels do diverge (sanity)
        m.atomizer.legacy_gen_labels = True
        with torch.no_grad():
            m.reset_state(reset_atomizer=True)
            pp = m.atomizer.encode_bytes(raw[:n0], reset=True)
            g2 = [m.forward_packet(p)["surface"]["byte_logits"].clone() for p in pp[:-1]]
            cur = pp[-1]
            for b in raw[n0:]:
                g2.append(m.forward_packet(cur)["surface"]["byte_logits"].clone())
                cur = m.atomizer.packet_from_payload(bytes([b]))
        self.assertGreater(max(float((a - b).abs().max()) for a, b in zip(tf, g2)), 0.0)

    def test_generate_legacy_flag_bit_identical_to_old_code(self) -> None:
        leg = _legacy()
        if leg is None:
            self.skipTest("legacy commit not available")
        OldModel, OldAtomizer = leg
        new = _build(field_leak=0.3)
        old = _build(cls=OldModel, atom_cls=OldAtomizer, field_leak=0.3)
        for x, y in zip(new.parameters(), old.parameters()):
            self.assertTrue(torch.equal(x, y))
        new.atomizer.legacy_gen_labels = True
        for prompt in ("Bonjour, comment ça va ?", "Il était une fois"):
            for seed in (0, 1, 2):
                torch.manual_seed(seed)
                a = new.generate_packets(prompt, max_packets=60, temperature=0.8, top_k=12,
                                         deterministic=False, max_length=60)
                torch.manual_seed(seed)
                b = old.generate_packets(prompt, max_packets=60, temperature=0.8, top_k=12,
                                         deterministic=False, max_length=60)
                self.assertEqual(a, b)


def _atom_grads_by_lag(model, n=24, lags=(1, 2, 3, 5, 8)):
    """||d loss_T / d atom.r_{T-L}|| on a train-mode tick sequence (no updates)."""
    model.train()
    model.surface.last_atom_dropout = 0.0
    stash = []
    orig = model.compiler.forward

    def hooked(features, atom_count=0):
        atom, op, conf = orig(features, atom_count)
        atom.r.retain_grad(); atom.E.retain_grad()
        stash.append(atom)
        return atom, op, conf

    model.compiler.forward = hooked
    try:
        pk = model.atomizer.encode(TEXT, reset=True)
        model.reset_state(reset_atomizer=False)
        for i in range(n - 1):
            model.forward_packet(pk[i])
        loss, _ = model.transition_loss(pk[n - 1], pk[n])
        loss.backward()
    finally:
        model.compiler.forward = orig
    out = {}
    for L in (0, *lags):
        a = stash[n - 1 - L]
        g = 0.0
        for t in (a.r, a.E):
            if t.grad is not None:
                g += float(t.grad.norm())
        out[L] = g
    return out


class FieldBpttTests(unittest.TestCase):
    def test_default_off(self) -> None:
        m = _build()
        self.assertEqual(m.field_bptt, 0)
        self.assertIsNone(m.field_leak_logit)
        self.assertEqual(m._bptt_checkpoint(), {})

    def test_off_bit_identical_to_legacy_commit(self) -> None:
        leg = _legacy()
        if leg is None:
            self.skipTest("legacy commit not available")
        OldModel, OldAtomizer = leg
        for kw in ({}, {"field_leak": 0.3}, {"field_leak_spectrum": (0.02, 0.3)}):
            _eq(self, _run(_build(**kw)), _run(_build(cls=OldModel, atom_cls=OldAtomizer, **kw)))

    def test_on_forward_values_identical(self) -> None:
        for kw in ({"field_leak": 0.3}, {"field_leak": 0.05}, {"field_leak_spectrum": (0.02, 0.3)}, {}):
            a, b = _build(**kw), _build(field_bptt=16, **kw)
            pk = a.atomizer.encode(TEXT, reset=True)
            outs = []
            for m in (a, b):
                m.train(); m.surface.last_atom_dropout = 0.0
                m.reset_state(reset_atomizer=False)
                torch.manual_seed(3)
                tr = []
                for cur, tgt in list(zip(pk[:-1], pk[1:]))[:30]:
                    loss, info = m.transition_loss(cur, tgt)
                    tr += [loss.detach().clone(), info["surface"]["byte_logits"].detach().clone(),
                           m.core.state.alpha.detach().clone()]
                outs.append(tr)
            _eq(self, outs[0], outs[1])
            self.assertIsNotNone(b._bptt_alpha)
            self.assertTrue(b._bptt_alpha.requires_grad)

    def test_eval_never_keeps_graph(self) -> None:
        m = _build(field_bptt=8, field_leak=0.05)
        pk = m.atomizer.encode(TEXT, reset=True)
        m.eval()
        with torch.no_grad():
            for p in pk[:10]:
                m.forward_packet(p)
        self.assertIsNone(m._bptt_alpha)
        m.train()
        with torch.no_grad():
            m.forward_packet(pk[10])
        self.assertIsNone(m._bptt_alpha)

    def test_gradient_reaches_past_atoms_only_when_on(self) -> None:
        off = _atom_grads_by_lag(_build(field_leak=0.05))
        self.assertGreater(off[0], 0.0)
        for L in (1, 2, 3, 5, 8):
            self.assertEqual(off[L], 0.0, f"lag {L} off")
        on = _atom_grads_by_lag(_build(field_leak=0.05, field_bptt=64))
        for L in (0, 1, 2, 3, 5, 8):
            self.assertGreater(on[L], 0.0, f"lag {L} on")

    def test_window_cut_stops_gradient(self) -> None:
        m = _build(field_leak=0.05, field_bptt=4)
        m.train(); m.surface.last_atom_dropout = 0.0
        stash = []
        orig = m.compiler.forward

        def hooked(features, atom_count=0):
            atom, op, conf = orig(features, atom_count)
            atom.r.retain_grad(); stash.append(atom)
            return atom, op, conf

        m.compiler.forward = hooked
        pk = m.atomizer.encode(TEXT, reset=True)
        m.reset_state(reset_atomizer=False)
        for i in range(6):
            m.forward_packet(pk[i])
        m.detach_field_graph()  # window boundary after tick 6
        for i in range(6, 9):
            m.forward_packet(pk[i])
        loss, _ = m.transition_loss(pk[9], pk[10])
        loss.backward()
        m.compiler.forward = orig
        self.assertIsNotNone(stash[7].r.grad)
        self.assertGreater(float(stash[7].r.grad.norm()), 0.0)
        for a in stash[:6]:
            self.assertTrue(a.r.grad is None or float(a.r.grad.norm()) == 0.0)

    def test_learnable_leak(self) -> None:
        # init equals the scalar leak (bit-identical forward trace, no BPTT, eval ticks)
        a, b = _build(field_leak=0.3), _build(field_leak=0.3, field_leak_learnable=True)
        lam = float(b.learnable_leak_rates().detach()[0])
        self.assertAlmostEqual(lam, 0.3, places=6)
        b2 = _build(field_leak=lam)
        pk = a.atomizer.encode(TEXT, reset=True)
        tr = []
        for m in (b2, b):
            m.eval(); m.reset_state(reset_atomizer=False)
            with torch.no_grad():
                tr.append([m.forward_packet(p)["surface"]["byte_logits"].clone() for p in pk[:20]])
        for x, y in zip(*tr):
            self.assertTrue(torch.allclose(x, y, atol=1e-6))
        # receives gradient under BPTT
        c = _build(field_leak=0.05, field_leak_learnable=True, field_bptt=32)
        c.train()
        c.reset_state(reset_atomizer=False)
        tot = 0
        for cur, tgt in list(zip(pk[:-1], pk[1:]))[:20]:
            loss, _ = c.transition_loss(cur, tgt); tot = tot + loss
        tot.backward()
        self.assertIsNotNone(c.field_leak_logit.grad)
        self.assertGreater(float(c.field_leak_logit.grad.abs().sum()), 0.0)
        with self.assertRaises(ValueError):
            _build(field_leak_spectrum=(0.02, 0.3), field_leak_learnable=True)

    def test_checkpoint_roundtrip(self) -> None:
        with tempfile.TemporaryDirectory() as t:
            m = _build(field_leak=0.05, field_bptt=64, field_leak_learnable=True)
            with torch.no_grad():
                m.field_leak_logit.add_(torch.linspace(-1, 1, 8).view(8, 1))
            p = Path(t) / "on.pt"
            m.save(p)
            f = _build()
            f.load(p)
            self.assertEqual(f.field_bptt, 64)
            self.assertTrue(torch.equal(f.field_leak_logit.detach(), m.field_leak_logit.detach()))
            # off ckpt: no new keys; loads as off into an "on" model
            p2 = Path(t) / "off.pt"
            _build().save(p2)
            blob = torch.load(p2, weights_only=False)
            self.assertNotIn("field_bptt", blob)
            self.assertNotIn("field_leak_logit", blob)
            g = _build(field_bptt=8, field_leak_learnable=True)
            g.load(p2)
            self.assertEqual(g.field_bptt, 0)
            self.assertIsNone(g.field_leak_logit)


class TrainerCliTests(unittest.TestCase):
    ENV = {"OMP_NUM_THREADS": "1", "MKL_NUM_THREADS": "1", "PATH": "/usr/bin:/bin"}
    BASE = ["--steps", "40", "--d-model", "8", "--n-modes", "8", "--n-atoms-max", "32",
            "--max-span-bytes", "1", "--field-obligatory-hard", "--no-enable-merge", "--no-payload-copy",
            "--last-atom-readout", "--allow-parallel-train", "--seed", "1", "--atom-flush-every", "16",
            "--episode-length", "25", "--efference-every", "10", "--field-loss-weight", "0.05",
            "--last-atom-dropout", "0.5", "--field-leak", "0.3"]

    def _train(self, tmp: Path, name: str, *extra: str, root: Path = ROOT) -> dict:
        data = tmp / "corpus.txt"
        if not data.exists():
            data.write_text(TEXT * 20, encoding="utf-8")
        cmd = [sys.executable, str(root / "tools" / "run_atom_native.py"), "--data", str(data),
               "--output-dir", str(tmp / name.replace(".pt", "_dir")), "--checkpoint-name", name,
               *self.BASE, *extra]
        subprocess.run(cmd, cwd=root, env={**self.ENV, "PYTHONPATH": str(root)}, check=True,
                       capture_output=True, text=True)
        return torch.load(tmp / name.replace(".pt", "_dir") / name, map_location="cpu", weights_only=False)

    def test_trainer_off_bit_identical_to_legacy_commit(self) -> None:
        with tempfile.TemporaryDirectory() as t:
            tmp = Path(t)
            legacy_root = tmp / "legacy_repo"
            try:
                legacy_root.mkdir()
                arch = subprocess.run(["git", "-C", str(ROOT), "archive", LEGACY_COMMIT, "src", "tools"],
                                      capture_output=True, check=True).stdout
                subprocess.run(["tar", "-x", "-C", str(legacy_root)], input=arch, check=True)
            except Exception:
                self.skipTest("legacy commit not available")
            # (a) defaults with efference off: bit-identical; (b) efference on +
            # --legacy-gen-labels: bit-identical (the fix changes efference commits by design)
            self._compare(self._train(tmp, "new.pt", "--efference-every", "0"),
                          self._train(tmp, "old.pt", "--efference-every", "0", root=legacy_root))
            self._compare(self._train(tmp, "new2.pt", "--efference-every", "2", "--legacy-gen-labels"),
                          self._train(tmp, "old2.pt", "--efference-every", "2", root=legacy_root))

    def _compare(self, new, old):
        if True:
            for part in ("compiler", "surface"):
                for k, v in old[part].items():
                    self.assertTrue(torch.equal(v, new[part][k]), f"{part}.{k}")
            self.assertTrue(torch.equal(old["core"]["state"]["alpha"], new["core"]["state"]["alpha"]))
            for k, v in old["core"]["dynamics"].items():
                self.assertTrue(torch.equal(v, new["core"]["dynamics"][k]))
            self.assertEqual(set(old.keys()), set(new.keys()))
            self.assertEqual(old["config"], new["config"])
            self.assertEqual([r["loss"] for r in old["training"]["trajectory"]],
                             [r["loss"] for r in new["training"]["trajectory"]])

    def test_cli_on_and_wins_after_resume(self) -> None:
        with tempfile.TemporaryDirectory() as t:
            tmp = Path(t)
            a = self._train(tmp, "a.pt", "--efference-every", "0")
            self.assertNotIn("field_bptt", a)
            b = self._train(tmp, "b.pt", "--efference-every", "0", "--field-bptt", "8",
                            "--field-leak", "0.05", "--field-leak-learnable")
            self.assertEqual(b["field_bptt"], 8)
            self.assertIn("field_leak_logit", b)
            # bptt ckpt resumed without the flags -> off
            c = self._train(tmp, "c.pt", "--efference-every", "0", "--resume", str(tmp / "b_dir" / "b.pt"))
            self.assertNotIn("field_bptt", c)
            self.assertNotIn("field_leak_logit", c)
            # plain ckpt resumed with the flag -> on; learnable leak kept from the ckpt when resumed with it
            d = self._train(tmp, "d.pt", "--efference-every", "0", "--resume", str(tmp / "a_dir" / "a.pt"),
                            "--field-bptt", "16")
            self.assertEqual(d["field_bptt"], 16)
            e = self._train(tmp, "e.pt", "--efference-every", "0", "--resume", str(tmp / "b_dir" / "b.pt"),
                            "--field-bptt", "8", "--field-leak-learnable", "--steps", "1")
            self.assertEqual(e["field_bptt"], 8)
            self.assertIn("field_leak_logit", e)
            # legacy gen labels flag runs (efference path) and is recorded only when on
            f = self._train(tmp, "f.pt", "--legacy-gen-labels")
            self.assertNotIn("field_bptt", f)


if __name__ == "__main__":
    unittest.main()

	"""Unit tests: byte-tick decode label fix + opt-in truncated BPTT through the field
	(docs/DECODE_FIX_AND_BPTT.md).

	Decode fix (automatic in byte-tick, escape flag ``legacy_gen_labels`` / --legacy-gen-labels):
	* byte-tick generated commits are bit-identical (label + features + atomizer state) to
	``encode_bytes`` of the same byte stream, in both pack modes;
	* the teacher-forced and generation paths give identical logits (max \|dlogit\| == 0);
	* ``legacy_gen_labels=True`` reproduces the pre-fix commit (commit f8a8918) exactly:
	labels, features, and a sampled generate_packets() byte string;
	* non-byte-tick (max_span > 1) commits are unchanged whatever the flag.

	Field BPTT (--field-bptt K, --field-leak-learnable):
	* off (default): multi-tick train+eval trace bit-identical to commit f8a8918
	(no leak, scalar leak 0.3, leak spectrum), trainer checkpoint bit-identical too;
	* on: forward values (field, logits, loss) bit-identical to off at equal weights;
	eval/no_grad ticks never keep a graph;
	* on: d loss_t / d atom_{t-L} != 0 for L = 1..K-1, and == 0 across a window cut;
	off: == 0 for every L >= 1 (the ROOT_CAUSE measurement);
	* learnable leak: init == scalar leak (bit-identical trace), receives gradient under BPTT;
	* checkpoint roundtrip; old checkpoints load as off; trainer CLI default off, on,
	and CLI wins after --resume in both directions.
	"""
	from __future__ import annotations

	import importlib
	import shutil
	import subprocess
	import sys
	import tempfile
	import unittest
	from pathlib import Path

	import torch

	from src.atom_native import AtomNativeModel
	from src.io.atomizer import Atomizer

	ROOT = Path(__file__).resolve().parents[1]
	LEGACY_COMMIT = "f8a8918"
	TEXT = "Bonjour, comment ça va ? Il était une fois un petit atome. Qui es-tu ?\nOui, d'accord... " * 3

	_LEG: dict = {}


	def _legacy():
	"""(AtomNativeModel, Atomizer) classes of the pre-change commit, or None."""
	if "mods" in _LEG:
	return _LEG["mods"]
	try:
	tmp = Path(tempfile.mkdtemp(prefix="legacy_dfb_"))
	pkg = tmp / "legacy_src_dfb"
	pkg.mkdir()
	files = subprocess.run(["git", "-C", str(ROOT), "ls-tree", "-r", "--name-only", LEGACY_COMMIT, "src"],
	capture_output=True, text=True, check=True).stdout.split()
	for f in files:
	if not f.endswith(".py"):
	continue
	code = subprocess.run(["git", "-C", str(ROOT), "show", f"{LEGACY_COMMIT}:{f}"],
	capture_output=True, text=True, check=True).stdout
	dst = pkg / Path(f).relative_to("src")
	dst.parent.mkdir(parents=True, exist_ok=True)
	dst.write_text(code, encoding="utf-8")
	for d in [pkg, [p for p in pkg.rglob("") if p.is_dir()]]:
	(d / "__init__.py").touch(exist_ok=True)
	sys.path.insert(0, str(tmp))
	try:
	an = importlib.import_module("legacy_src_dfb.atom_native")
	at = importlib.import_module("legacy_src_dfb.io.atomizer")
	finally:
	sys.path.remove(str(tmp))
	_LEG["mods"] = (an.AtomNativeModel, at.Atomizer)
	except Exception: # pragma: no cover
	_LEG["mods"] = None
	return _LEG["mods"]


	def _build(cls=AtomNativeModel, atom_cls=Atomizer, **kw):
	torch.manual_seed(0)
	return cls(
	d_model=8, n_modes=8, n_atoms_max=32, max_payload_bytes=1,
	atomizer=atom_cls(max_span_bytes=1), field_max_rms=3.0,
	field_obligatory_hard=True, enable_merge=False, last_atom_readout=True,
	field_loss_weight=0.05, **kw,
	)


	def _run(model, n_train=40, n_eval=20):
	opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=1e-3)
	packets = model.atomizer.encode(TEXT, reset=True)
	model.reset_state(reset_atomizer=False)
	model.train()
	trace = []
	torch.manual_seed(5)
	for cur, tgt in list(zip(packets[:-1], packets[1:]))[:n_train]:
	opt.zero_grad(set_to_none=True)
	loss, _ = model.transition_loss(cur, tgt)
	loss.backward()
	opt.step()
	model.stabilize_dynamics_parameters()
	trace.append(loss.detach().clone())
	trace.append(model.core.state.alpha.detach().clone())
	model.eval()
	with torch.no_grad():
	for cur, tgt in list(zip(packets[:-1], packets[1:]))[n_train:n_train + n_eval]:
	loss, _ = model.transition_loss(cur, tgt)
	trace.append(loss.detach().clone())
	trace.append(model.core.state.alpha.detach().clone())
	trace.extend(p.detach().clone() for p in model.parameters())
	return trace


	def _eq(tc, a, b):
	tc.assertEqual(len(a), len(b))
	for x, y in zip(a, b):
	tc.assertTrue(torch.equal(x, y))


	class DecodeLabelFixTests(unittest.TestCase):
	def test_byte_tick_commit_equals_training_encode(self) -> None:
	raw = TEXT.encode("utf-8")
	for mode in ("linguistic", "eager"):
	tr = Atomizer(max_span_bytes=1, pack_mode=mode).encode_bytes(raw, reset=True, flush=False)
	gen_at = Atomizer(max_span_bytes=1, pack_mode=mode)
	gen = [gen_at.packet_from_payload(bytes([b])) for b in raw]
	self.assertEqual(len(tr), len(gen))
	for a, b in zip(tr, gen):
	self.assertEqual((a.payload, a.start, a.boundary, a.level), (b.payload, b.start, b.boundary, b.level))
	self.assertTrue(torch.equal(a.features, b.features))
	self.assertEqual({p.boundary for p in gen} - {"max_span", "newline", "punctuation", "whitespace"}, set())
	if mode == "linguistic":
	self.assertEqual({p.boundary for p in gen}, {"max_span", "newline"})

	def test_legacy_flag_reproduces_old_labels_exactly(self) -> None:
	leg = _legacy()
	if leg is None:
	self.skipTest("legacy commit not available")
	_, OldAtomizer = leg
	raw = TEXT.encode("utf-8")
	for mode in ("linguistic", "eager"):
	new = Atomizer(max_span_bytes=1, pack_mode=mode); new.legacy_gen_labels = True
	old = OldAtomizer(max_span_bytes=1, pack_mode=mode)
	for b in raw:
	a, o = new.packet_from_payload(bytes([b])), old.packet_from_payload(bytes([b]))
	self.assertEqual((a.boundary, a.start, a.level), (o.boundary, o.start, o.level))
	self.assertTrue(torch.equal(a.features, o.features))
	# the fix does change labels (else the test above would prove nothing)
	fixed = Atomizer(max_span_bytes=1)
	self.assertEqual(fixed.packet_from_payload(b" ").boundary, "max_span")
	legacy = Atomizer(max_span_bytes=1); legacy.legacy_gen_labels = True
	self.assertEqual(legacy.packet_from_payload(b" ").boundary, "whitespace")

	def test_non_byte_tick_unchanged(self) -> None:
	leg = _legacy()
	if leg is None:
	self.skipTest("legacy commit not available")
	_, OldAtomizer = leg
	for flag in (False, True):
	new = Atomizer(max_span_bytes=8); new.legacy_gen_labels = flag
	old = OldAtomizer(max_span_bytes=8)
	for pl in (b"abc ", b"Oui,", b"x", b" ", b"\n", b"mot"):
	a, o = new.packet_from_payload(pl), old.packet_from_payload(pl)
	self.assertEqual(a.boundary, o.boundary)
	self.assertTrue(torch.equal(a.features, o.features))

	def test_teacher_forced_equals_generation_path(self) -> None:
	m = _build(field_leak=0.3)
	m.eval()
	text = "Utilisateur: Qui es-tu ?\nAssistant: Je suis Atom, un petit modèle, d'accord ?\n"
	raw = text.encode()
	with torch.no_grad():
	m.reset_state(reset_atomizer=True)
	pk = m.atomizer.encode_bytes(raw, reset=True)
	tf = [m.forward_packet(p)["surface"]["byte_logits"].clone() for p in pk[:-1]]
	m.reset_state(reset_atomizer=True)
	n0 = 30
	pp = m.atomizer.encode_bytes(raw[:n0], reset=True)
	gen = [m.forward_packet(p)["surface"]["byte_logits"].clone() for p in pp[:-1]]
	cur = pp[-1]
	for b in raw[n0:]:
	gen.append(m.forward_packet(cur)["surface"]["byte_logits"].clone())
	cur = m.atomizer.packet_from_payload(bytes([b]))
	self.assertEqual(len(tf), len(gen))
	self.assertEqual(max(float((a - b).abs().max()) for a, b in zip(tf, gen)), 0.0)
	# legacy labels do diverge (sanity)
	m.atomizer.legacy_gen_labels = True
	with torch.no_grad():
	m.reset_state(reset_atomizer=True)
	pp = m.atomizer.encode_bytes(raw[:n0], reset=True)
	g2 = [m.forward_packet(p)["surface"]["byte_logits"].clone() for p in pp[:-1]]
	cur = pp[-1]
	for b in raw[n0:]:
	g2.append(m.forward_packet(cur)["surface"]["byte_logits"].clone())
	cur = m.atomizer.packet_from_payload(bytes([b]))
	self.assertGreater(max(float((a - b).abs().max()) for a, b in zip(tf, g2)), 0.0)

	def test_generate_legacy_flag_bit_identical_to_old_code(self) -> None:
	leg = _legacy()
	if leg is None:
	self.skipTest("legacy commit not available")
	OldModel, OldAtomizer = leg
	new = _build(field_leak=0.3)
	old = _build(cls=OldModel, atom_cls=OldAtomizer, field_leak=0.3)
	for x, y in zip(new.parameters(), old.parameters()):
	self.assertTrue(torch.equal(x, y))
	new.atomizer.legacy_gen_labels = True
	for prompt in ("Bonjour, comment ça va ?", "Il était une fois"):
	for seed in (0, 1, 2):
	torch.manual_seed(seed)
	a = new.generate_packets(prompt, max_packets=60, temperature=0.8, top_k=12,
	deterministic=False, max_length=60)
	torch.manual_seed(seed)
	b = old.generate_packets(prompt, max_packets=60, temperature=0.8, top_k=12,
	deterministic=False, max_length=60)
	self.assertEqual(a, b)


	def _atom_grads_by_lag(model, n=24, lags=(1, 2, 3, 5, 8)):
	"""\|\|d loss_T / d atom.r_{T-L}\|\| on a train-mode tick sequence (no updates)."""
	model.train()
	model.surface.last_atom_dropout = 0.0
	stash = []
	orig = model.compiler.forward

	def hooked(features, atom_count=0):
	atom, op, conf = orig(features, atom_count)
	atom.r.retain_grad(); atom.E.retain_grad()
	stash.append(atom)
	return atom, op, conf

	model.compiler.forward = hooked
	try:
	pk = model.atomizer.encode(TEXT, reset=True)
	model.reset_state(reset_atomizer=False)
	for i in range(n - 1):
	model.forward_packet(pk[i])
	loss, _ = model.transition_loss(pk[n - 1], pk[n])
	loss.backward()
	finally:
	model.compiler.forward = orig
	out = {}
	for L in (0, *lags):
	a = stash[n - 1 - L]
	g = 0.0
	for t in (a.r, a.E):
	if t.grad is not None:
	g += float(t.grad.norm())
	out[L] = g
	return out


	class FieldBpttTests(unittest.TestCase):
	def test_default_off(self) -> None:
	m = _build()
	self.assertEqual(m.field_bptt, 0)
	self.assertIsNone(m.field_leak_logit)
	self.assertEqual(m._bptt_checkpoint(), {})

	def test_off_bit_identical_to_legacy_commit(self) -> None:
	leg = _legacy()
	if leg is None:
	self.skipTest("legacy commit not available")
	OldModel, OldAtomizer = leg
	for kw in ({}, {"field_leak": 0.3}, {"field_leak_spectrum": (0.02, 0.3)}):
	_eq(self, _run(_build(kw)), _run(_build(cls=OldModel, atom_cls=OldAtomizer, kw)))

	def test_on_forward_values_identical(self) -> None:
	for kw in ({"field_leak": 0.3}, {"field_leak": 0.05}, {"field_leak_spectrum": (0.02, 0.3)}, {}):
	a, b = _build(kw), _build(field_bptt=16, kw)
	pk = a.atomizer.encode(TEXT, reset=True)
	outs = []
	for m in (a, b):
	m.train(); m.surface.last_atom_dropout = 0.0
	m.reset_state(reset_atomizer=False)
	torch.manual_seed(3)
	tr = []
	for cur, tgt in list(zip(pk[:-1], pk[1:]))[:30]:
	loss, info = m.transition_loss(cur, tgt)
	tr += [loss.detach().clone(), info["surface"]["byte_logits"].detach().clone(),
	m.core.state.alpha.detach().clone()]
	outs.append(tr)
	_eq(self, outs[0], outs[1])
	self.assertIsNotNone(b._bptt_alpha)
	self.assertTrue(b._bptt_alpha.requires_grad)

	def test_eval_never_keeps_graph(self) -> None:
	m = _build(field_bptt=8, field_leak=0.05)
	pk = m.atomizer.encode(TEXT, reset=True)
	m.eval()
	with torch.no_grad():
	for p in pk[:10]:
	m.forward_packet(p)
	self.assertIsNone(m._bptt_alpha)
	m.train()
	with torch.no_grad():
	m.forward_packet(pk[10])
	self.assertIsNone(m._bptt_alpha)

	def test_gradient_reaches_past_atoms_only_when_on(self) -> None:
	off = _atom_grads_by_lag(_build(field_leak=0.05))
	self.assertGreater(off[0], 0.0)
	for L in (1, 2, 3, 5, 8):
	self.assertEqual(off[L], 0.0, f"lag {L} off")
	on = _atom_grads_by_lag(_build(field_leak=0.05, field_bptt=64))
	for L in (0, 1, 2, 3, 5, 8):
	self.assertGreater(on[L], 0.0, f"lag {L} on")

	def test_window_cut_stops_gradient(self) -> None:
	m = _build(field_leak=0.05, field_bptt=4)
	m.train(); m.surface.last_atom_dropout = 0.0
	stash = []
	orig = m.compiler.forward

	def hooked(features, atom_count=0):
	atom, op, conf = orig(features, atom_count)
	atom.r.retain_grad(); stash.append(atom)
	return atom, op, conf

	m.compiler.forward = hooked
	pk = m.atomizer.encode(TEXT, reset=True)
	m.reset_state(reset_atomizer=False)
	for i in range(6):
	m.forward_packet(pk[i])
	m.detach_field_graph() # window boundary after tick 6
	for i in range(6, 9):
	m.forward_packet(pk[i])
	loss, _ = m.transition_loss(pk[9], pk[10])
	loss.backward()
	m.compiler.forward = orig
	self.assertIsNotNone(stash[7].r.grad)
	self.assertGreater(float(stash[7].r.grad.norm()), 0.0)
	for a in stash[:6]:
	self.assertTrue(a.r.grad is None or float(a.r.grad.norm()) == 0.0)

	def test_learnable_leak(self) -> None:
	# init equals the scalar leak (bit-identical forward trace, no BPTT, eval ticks)
	a, b = _build(field_leak=0.3), _build(field_leak=0.3, field_leak_learnable=True)
	lam = float(b.learnable_leak_rates().detach()[0])
	self.assertAlmostEqual(lam, 0.3, places=6)
	b2 = _build(field_leak=lam)
	pk = a.atomizer.encode(TEXT, reset=True)
	tr = []
	for m in (b2, b):
	m.eval(); m.reset_state(reset_atomizer=False)
	with torch.no_grad():
	tr.append([m.forward_packet(p)["surface"]["byte_logits"].clone() for p in pk[:20]])
	for x, y in zip(*tr):
	self.assertTrue(torch.allclose(x, y, atol=1e-6))
	# receives gradient under BPTT
	c = _build(field_leak=0.05, field_leak_learnable=True, field_bptt=32)
	c.train()
	c.reset_state(reset_atomizer=False)
	tot = 0
	for cur, tgt in list(zip(pk[:-1], pk[1:]))[:20]:
	loss, _ = c.transition_loss(cur, tgt); tot = tot + loss
	tot.backward()
	self.assertIsNotNone(c.field_leak_logit.grad)
	self.assertGreater(float(c.field_leak_logit.grad.abs().sum()), 0.0)
	with self.assertRaises(ValueError):
	_build(field_leak_spectrum=(0.02, 0.3), field_leak_learnable=True)

	def test_checkpoint_roundtrip(self) -> None:
	with tempfile.TemporaryDirectory() as t:
	m = _build(field_leak=0.05, field_bptt=64, field_leak_learnable=True)
	with torch.no_grad():
	m.field_leak_logit.add_(torch.linspace(-1, 1, 8).view(8, 1))
	p = Path(t) / "on.pt"
	m.save(p)
	f = _build()
	f.load(p)
	self.assertEqual(f.field_bptt, 64)
	self.assertTrue(torch.equal(f.field_leak_logit.detach(), m.field_leak_logit.detach()))
	# off ckpt: no new keys; loads as off into an "on" model
	p2 = Path(t) / "off.pt"
	_build().save(p2)
	blob = torch.load(p2, weights_only=False)
	self.assertNotIn("field_bptt", blob)
	self.assertNotIn("field_leak_logit", blob)
	g = _build(field_bptt=8, field_leak_learnable=True)
	g.load(p2)
	self.assertEqual(g.field_bptt, 0)
	self.assertIsNone(g.field_leak_logit)


	class TrainerCliTests(unittest.TestCase):
	ENV = {"OMP_NUM_THREADS": "1", "MKL_NUM_THREADS": "1", "PATH": "/usr/bin:/bin"}
	BASE = ["--steps", "40", "--d-model", "8", "--n-modes", "8", "--n-atoms-max", "32",
	"--max-span-bytes", "1", "--field-obligatory-hard", "--no-enable-merge", "--no-payload-copy",
	"--last-atom-readout", "--allow-parallel-train", "--seed", "1", "--atom-flush-every", "16",
	"--episode-length", "25", "--efference-every", "10", "--field-loss-weight", "0.05",
	"--last-atom-dropout", "0.5", "--field-leak", "0.3"]

	def _train(self, tmp: Path, name: str, *extra: str, root: Path = ROOT) -> dict:
	data = tmp / "corpus.txt"
	if not data.exists():
	data.write_text(TEXT * 20, encoding="utf-8")
	cmd = [sys.executable, str(root / "tools" / "run_atom_native.py"), "--data", str(data),
	"--output-dir", str(tmp / name.replace(".pt", "_dir")), "--checkpoint-name", name,
	self.BASE, extra]
	subprocess.run(cmd, cwd=root, env={**self.ENV, "PYTHONPATH": str(root)}, check=True,
	capture_output=True, text=True)
	return torch.load(tmp / name.replace(".pt", "_dir") / name, map_location="cpu", weights_only=False)

	def test_trainer_off_bit_identical_to_legacy_commit(self) -> None:
	with tempfile.TemporaryDirectory() as t:
	tmp = Path(t)
	legacy_root = tmp / "legacy_repo"
	try:
	legacy_root.mkdir()
	arch = subprocess.run(["git", "-C", str(ROOT), "archive", LEGACY_COMMIT, "src", "tools"],
	capture_output=True, check=True).stdout
	subprocess.run(["tar", "-x", "-C", str(legacy_root)], input=arch, check=True)
	except Exception:
	self.skipTest("legacy commit not available")
	# (a) defaults with efference off: bit-identical; (b) efference on +
	# --legacy-gen-labels: bit-identical (the fix changes efference commits by design)
	self._compare(self._train(tmp, "new.pt", "--efference-every", "0"),
	self._train(tmp, "old.pt", "--efference-every", "0", root=legacy_root))
	self._compare(self._train(tmp, "new2.pt", "--efference-every", "2", "--legacy-gen-labels"),
	self._train(tmp, "old2.pt", "--efference-every", "2", root=legacy_root))

	def _compare(self, new, old):
	if True:
	for part in ("compiler", "surface"):
	for k, v in old[part].items():
	self.assertTrue(torch.equal(v, new[part][k]), f"{part}.{k}")
	self.assertTrue(torch.equal(old["core"]["state"]["alpha"], new["core"]["state"]["alpha"]))
	for k, v in old["core"]["dynamics"].items():
	self.assertTrue(torch.equal(v, new["core"]["dynamics"][k]))
	self.assertEqual(set(old.keys()), set(new.keys()))
	self.assertEqual(old["config"], new["config"])
	self.assertEqual([r["loss"] for r in old["training"]["trajectory"]],
	[r["loss"] for r in new["training"]["trajectory"]])

	def test_cli_on_and_wins_after_resume(self) -> None:
	with tempfile.TemporaryDirectory() as t:
	tmp = Path(t)
	a = self._train(tmp, "a.pt", "--efference-every", "0")
	self.assertNotIn("field_bptt", a)
	b = self._train(tmp, "b.pt", "--efference-every", "0", "--field-bptt", "8",
	"--field-leak", "0.05", "--field-leak-learnable")
	self.assertEqual(b["field_bptt"], 8)
	self.assertIn("field_leak_logit", b)
	# bptt ckpt resumed without the flags -> off
	c = self._train(tmp, "c.pt", "--efference-every", "0", "--resume", str(tmp / "b_dir" / "b.pt"))
	self.assertNotIn("field_bptt", c)
	self.assertNotIn("field_leak_logit", c)
	# plain ckpt resumed with the flag -> on; learnable leak kept from the ckpt when resumed with it
	d = self._train(tmp, "d.pt", "--efference-every", "0", "--resume", str(tmp / "a_dir" / "a.pt"),
	"--field-bptt", "16")
	self.assertEqual(d["field_bptt"], 16)
	e = self._train(tmp, "e.pt", "--efference-every", "0", "--resume", str(tmp / "b_dir" / "b.pt"),
	"--field-bptt", "8", "--field-leak-learnable", "--steps", "1")
	self.assertEqual(e["field_bptt"], 8)
	self.assertIn("field_leak_logit", e)
	# legacy gen labels flag runs (efference path) and is recorded only when on
	f = self._train(tmp, "f.pt", "--legacy-gen-labels")
	self.assertNotIn("field_bptt", f)


	if __name__ == "__main__":
	unittest.main()