summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
authorYurenHao0426 <Blackhao0426@gmail.com>2026-07-30 16:34:30 -0500
committerYurenHao0426 <Blackhao0426@gmail.com>2026-07-30 16:34:30 -0500
commit716205eba862076c1d0fbbff016601272e4614f2 (patch)
treedc8c31dcfefa49d53c98d305fe35419da76e58c9
parent906d17e896329e4d2dd8d7f728e4d5a234afe0e4 (diff)
experiment: package portable A6000 crossover
-rw-r--r--experiments/analyze_resnet_crossover_r2.py29
-rw-r--r--experiments/analyze_transformer_crossover_t2.py30
-rwxr-xr-xexperiments/bootstrap_plain_cnn.sh35
-rw-r--r--experiments/crossover_hardware.py124
-rwxr-xr-xexperiments/crossover_hardware_smoke.py29
-rwxr-xr-xexperiments/prepare_shakespeare_char.py77
-rw-r--r--experiments/resnet_crossover_r2.py62
-rwxr-xr-xexperiments/run_a6000_matrix.py182
-rw-r--r--experiments/transformer_crossover_native.py5
-rw-r--r--experiments/transformer_crossover_t2.py62
-rw-r--r--external/README.md11
-rw-r--r--external/dualprop_patches/0001-crossover-separate-author-diagnostics-from-timed-tra.patch198
-rw-r--r--external/dualprop_patches/0002-crossover-add-audited-ordinary-FA-and-DFA-rules.patch453
-rw-r--r--external/dualprop_patches/0003-fix-norm-preserve-direct-feedback-maps.patch43
-rw-r--r--external/dualprop_patches/0004-crossover-add-architecture-compatible-PEPITA.patch242
-rw-r--r--external/dualprop_patches/0005-crossover-add-greedy-Forward-Forward-runner.patch422
-rw-r--r--external/dualprop_patches/0006-crossover-add-two-phase-equilibrium-propagation.patch377
-rw-r--r--external/dualprop_patches/0007-crossover-add-reciprocal-clean-KP-training.patch321
-rw-r--r--external/dualprop_patches/0008-crossover-add-dynamic-innovation-on-reciprocal-KP.patch430
-rw-r--r--external/dualprop_patches/0009-crossover-isolate-formal-validation-from-test.patch135
-rw-r--r--external/dualprop_patches/0010-crossover-add-paper-PEPITA-learning-schedule.patch110
-rw-r--r--external/dualprop_patches/0011-experiment-register-plain-CNN-P1-selector-grid.patch247
-rw-r--r--external/dualprop_patches/0012-fix-map-crossover-BP-to-author-CLI.patch24
-rw-r--r--external/dualprop_patches/0013-analysis-audit-plain-CNN-P1-selector.patch260
-rw-r--r--external/dualprop_patches/0014-experiment-freeze-plain-CNN-P2-registry.patch228
-rw-r--r--external/dualprop_patches/0015-experiment-audit-P2-execution-and-resource-use.patch341
-rw-r--r--external/dualprop_patches/0016-experiment-lock-P2-source-and-cell-registry.patch70
-rw-r--r--external/dualprop_patches/0017-analysis-freeze-complete-P2-audit-gate.patch248
-rw-r--r--external/dualprop_patches/0018-fix-separate-P2-source-from-cell-hardware.patch24
-rw-r--r--external/dualprop_patches/0019-experiment-add-portable-A6000-P2-profile.patch217
-rw-r--r--sdil/data.py3
31 files changed, 4945 insertions, 94 deletions
diff --git a/experiments/analyze_resnet_crossover_r2.py b/experiments/analyze_resnet_crossover_r2.py
index 3c2bf46..ce89165 100644
--- a/experiments/analyze_resnet_crossover_r2.py
+++ b/experiments/analyze_resnet_crossover_r2.py
@@ -18,6 +18,7 @@ from experiments.resnet_crossover_r2 import (
registry_sha256,
selector_report,
)
+from experiments.crossover_hardware import assert_hardware_report
def sha256(path):
@@ -216,7 +217,9 @@ def work_metrics(record, job, history):
}
-def audit_job(job, expected_source, expected_selector):
+def audit_job(
+ job, expected_source, expected_selector, hardware_policy
+):
manifest_path = job["output"] + ".manifest.json"
if not os.path.isfile(manifest_path):
raise AssertionError(f"missing R2 manifest: {job['experiment_name']}")
@@ -240,8 +243,7 @@ def audit_job(job, expected_source, expected_selector):
assert manifest["source"] == expected_source
assert manifest["selector"] == expected_selector
hardware = manifest["hardware_lock"]
- assert hardware["physical_gpu_index"] in (5, 7)
- assert hardware["physical_gpu_uuid"]
+ assert_hardware_report(hardware, hardware_policy)
common = {
"cell_id": f"resnet{job['depth']}::{job['method']}",
"method": job["method"],
@@ -331,8 +333,14 @@ def audit_job(job, expected_source, expected_selector):
expected_ordinary *= job["depth"]
assert work["ordinary_training_examples"] == expected_ordinary
hardware_record = record["hardware"]
- assert hardware_record["cuda_visible_devices"] in ("5", "7")
- assert hardware_record["cuda_device_name"] == "NVIDIA GeForce GTX 1080"
+ assert (
+ hardware_record["cuda_visible_devices"]
+ == hardware["cuda_visible_devices"]
+ )
+ assert (
+ hardware_record["cuda_device_name"]
+ == hardware["physical_gpu_name"]
+ )
parameter_count = architecture.get("forward_parameter_count")
if parameter_count is None:
parameter_count = architecture["forward_parameters"]
@@ -370,9 +378,15 @@ def main():
assert launch["selector"] == selector
assert launch["registry_sha256"] == registry_sha256(jobs)
assert launch["num_jobs"] == 27
- assert launch["allowed_physical_gpus"] == [5, 7]
+ hardware_policy = launch["hardware_policy"]
+ assert (
+ launch["allowed_physical_gpus"]
+ == hardware_policy["allowed_physical_gpu_indices"]
+ )
source = launch["source"]
- records = [audit_job(job, source, selector) for job in jobs]
+ records = [
+ audit_job(job, source, selector, hardware_policy) for job in jobs
+ ]
assert len(records) == 27
assert len({record["cell_id"] for record in records}) == 27
assert {
@@ -401,6 +415,7 @@ def main():
"test_policy": "none",
"source": source,
"selector": selector,
+ "hardware_policy": hardware_policy,
"launch_lock": {
"path": os.path.relpath(launch_path, ROOT),
"sha256": sha256(launch_path),
diff --git a/experiments/analyze_transformer_crossover_t2.py b/experiments/analyze_transformer_crossover_t2.py
index 5eeb226..74c3ef6 100644
--- a/experiments/analyze_transformer_crossover_t2.py
+++ b/experiments/analyze_transformer_crossover_t2.py
@@ -18,6 +18,7 @@ from experiments.transformer_crossover_t2 import (
selector_report,
t2_jobs,
)
+from experiments.crossover_hardware import assert_hardware_report
TRAIN_HASH = (
@@ -79,7 +80,7 @@ def expected_work(record, job):
return work
-def audit_completed(job, manifest, source):
+def audit_completed(job, manifest, source, hardware_policy):
assert manifest["output_exists"] is True
assert manifest["output_sha256"] == sha256(job["output"])
record = read_json(job["output"])
@@ -164,8 +165,13 @@ def audit_completed(job, manifest, source):
)
assert record["work"]["logical_task_loss_queries"] == 0
hardware = record["hardware"]
- assert hardware["cuda_visible_devices"] in ("5", "7")
- assert hardware["device_name"] == "NVIDIA GeForce GTX 1080"
+ locked_hardware = manifest["hardware_lock"]
+ assert_hardware_report(locked_hardware, hardware_policy)
+ assert (
+ hardware["cuda_visible_devices"]
+ == locked_hardware["cuda_visible_devices"]
+ )
+ assert hardware["device_name"] == locked_hardware["physical_gpu_name"]
assert hardware["peak_memory_allocated_bytes"] is not None
assert hardware["peak_memory_reserved_bytes"] is not None
return {
@@ -189,7 +195,7 @@ def audit_completed(job, manifest, source):
}
-def audit_job(job, source, selector):
+def audit_job(job, source, selector, hardware_policy):
manifest_path = job["output"] + ".manifest.json"
if not os.path.isfile(manifest_path):
raise AssertionError(f"missing T2 manifest: {job['experiment_name']}")
@@ -211,8 +217,7 @@ def audit_job(job, source, selector):
assert manifest["source"] == source
assert manifest["selector"] == selector
hardware = manifest["hardware_lock"]
- assert hardware["physical_gpu_index"] in (5, 7)
- assert hardware["physical_gpu_uuid"]
+ assert_hardware_report(hardware, hardware_policy)
common = {
"cell_id": f"transformer{job['depth']}::{job['method']}",
"method": job["method"],
@@ -228,7 +233,7 @@ def audit_job(job, source, selector):
if manifest["status"] == "completed":
return {
**common,
- **audit_completed(job, manifest, source),
+ **audit_completed(job, manifest, source, hardware_policy),
}
assert manifest["status"] in {
"timeout", "nonzero_exit", "missing_output"
@@ -271,9 +276,15 @@ def main():
assert launch["selector"] == selector
assert launch["registry_sha256"] == registry_sha256(jobs)
assert launch["num_jobs"] == 27
- assert launch["allowed_physical_gpus"] == [5, 7]
+ hardware_policy = launch["hardware_policy"]
+ assert (
+ launch["allowed_physical_gpus"]
+ == hardware_policy["allowed_physical_gpu_indices"]
+ )
source = launch["source"]
- records = [audit_job(job, source, selector) for job in jobs]
+ records = [
+ audit_job(job, source, selector, hardware_policy) for job in jobs
+ ]
assert len(records) == 27
assert len({record["cell_id"] for record in records}) == 27
assert {
@@ -294,6 +305,7 @@ def main():
"test_policy": "none",
"source": source,
"selector": selector,
+ "hardware_policy": hardware_policy,
"launch_lock": {
"path": os.path.relpath(launch_path, ROOT),
"sha256": sha256(launch_path),
diff --git a/experiments/bootstrap_plain_cnn.sh b/experiments/bootstrap_plain_cnn.sh
new file mode 100755
index 0000000..0d6145f
--- /dev/null
+++ b/experiments/bootstrap_plain_cnn.sh
@@ -0,0 +1,35 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+if [[ "$#" -ne 1 ]]; then
+ echo "usage: $0 /absolute/path/to/dualprop-sdil" >&2
+ exit 2
+fi
+
+target="$1"
+if [[ "$target" != /* ]]; then
+ echo "target must be an absolute path" >&2
+ exit 2
+fi
+if [[ -e "$target" ]]; then
+ echo "refusing to overwrite existing target: $target" >&2
+ exit 2
+fi
+
+main_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
+patch_root="$main_root/external/dualprop_patches"
+base_revision="7b2595b34421e1483a721dbfdeff8cdabda3a1ff"
+
+if [[ "$(find "$patch_root" -maxdepth 1 -name '*.patch' | wc -l)" -ne 19 ]]; then
+ echo "expected 19 frozen Dual Propagation patches" >&2
+ exit 1
+fi
+
+git clone https://github.com/Rasmuskh/dualprop_icml_2024.git "$target"
+git -C "$target" checkout -b sdil-a6000 "$base_revision"
+git -C "$target" config user.name "SDIL replication runner"
+git -C "$target" config user.email "sdil-replication@invalid.example"
+git -C "$target" am --committer-date-is-author-date "$patch_root"/*.patch
+
+echo "plain-CNN runner ready at $target"
+git -C "$target" log -1 --oneline
diff --git a/experiments/crossover_hardware.py b/experiments/crossover_hardware.py
new file mode 100644
index 0000000..6ace655
--- /dev/null
+++ b/experiments/crossover_hardware.py
@@ -0,0 +1,124 @@
+#!/usr/bin/env python3
+"""Hardware profiles shared by portable crossover launchers and audits."""
+import os
+import subprocess
+
+
+PROFILES = {
+ "timan107_gtx1080": {
+ "allowed_physical_gpu_indices": [5, 7],
+ "allowed_device_names": ["NVIDIA GeForce GTX 1080"],
+ },
+ "a6000": {
+ "allowed_physical_gpu_indices": None,
+ "allowed_device_names": ["NVIDIA RTX A6000"],
+ },
+}
+DEFAULT_PROFILE = "timan107_gtx1080"
+
+
+def profile_choices():
+ return tuple(PROFILES)
+
+
+def policy(profile, physical_gpu_index=None):
+ if profile not in PROFILES:
+ raise ValueError(f"unknown crossover hardware profile: {profile}")
+ configured = PROFILES[profile]
+ allowed = configured["allowed_physical_gpu_indices"]
+ if allowed is None:
+ if physical_gpu_index is None:
+ raise ValueError(
+ f"{profile} policy requires a resolved physical GPU index"
+ )
+ allowed = [int(physical_gpu_index)]
+ return {
+ "profile": profile,
+ "allowed_physical_gpu_indices": list(allowed),
+ "allowed_device_names": list(configured["allowed_device_names"]),
+ "single_gpu_per_process": True,
+ }
+
+
+def _query_physical_gpus():
+ query = subprocess.run(
+ [
+ "nvidia-smi",
+ "--query-gpu=index,uuid,name,memory.total",
+ "--format=csv,noheader,nounits",
+ ],
+ check=True,
+ capture_output=True,
+ text=True,
+ ).stdout.splitlines()
+ rows = {}
+ for line in query:
+ index, uuid, name, memory = [
+ value.strip() for value in line.split(",", 3)
+ ]
+ rows[index] = {
+ "uuid": uuid,
+ "name": name,
+ "total_memory_mib": int(memory),
+ }
+ return rows
+
+
+def physical_gpu_report(profile, dry_run=False):
+ visible = os.environ.get("CUDA_VISIBLE_DEVICES")
+ if dry_run:
+ return {
+ "hardware_profile": profile,
+ "cuda_visible_devices": visible,
+ "physical_gpu_index": None,
+ "physical_gpu_uuid": None,
+ "physical_gpu_name": None,
+ "total_memory_mib": None,
+ }
+ if visible is None or not visible.isdigit():
+ raise RuntimeError(
+ "formal crossover jobs require CUDA_VISIBLE_DEVICES to contain "
+ "one physical numeric GPU index"
+ )
+ physical_index = int(visible)
+ hardware_policy = policy(profile, physical_index)
+ if physical_index not in hardware_policy["allowed_physical_gpu_indices"]:
+ raise RuntimeError(
+ f"physical GPU {physical_index} is not allowed by {profile}"
+ )
+ rows = _query_physical_gpus()
+ if visible not in rows:
+ raise RuntimeError(f"physical GPU {visible} not found")
+ row = rows[visible]
+ if row["name"] not in hardware_policy["allowed_device_names"]:
+ raise RuntimeError(
+ f"{profile} requires one of "
+ f"{hardware_policy['allowed_device_names']}, found {row['name']}"
+ )
+ return {
+ "hardware_profile": profile,
+ "cuda_visible_devices": visible,
+ "physical_gpu_index": physical_index,
+ "physical_gpu_uuid": row["uuid"],
+ "physical_gpu_name": row["name"],
+ "total_memory_mib": row["total_memory_mib"],
+ }
+
+
+def policy_for_report(profile, report):
+ return policy(profile, report["physical_gpu_index"])
+
+
+def assert_hardware_report(report, hardware_policy):
+ assert report["hardware_profile"] == hardware_policy["profile"]
+ assert (
+ report["physical_gpu_index"]
+ in hardware_policy["allowed_physical_gpu_indices"]
+ )
+ assert report["physical_gpu_uuid"]
+ assert report["physical_gpu_name"] in hardware_policy[
+ "allowed_device_names"
+ ]
+ assert report["cuda_visible_devices"] == str(
+ report["physical_gpu_index"]
+ )
diff --git a/experiments/crossover_hardware_smoke.py b/experiments/crossover_hardware_smoke.py
new file mode 100755
index 0000000..89c11c2
--- /dev/null
+++ b/experiments/crossover_hardware_smoke.py
@@ -0,0 +1,29 @@
+#!/usr/bin/env python3
+"""Dependency-free checks for frozen crossover hardware profiles."""
+from crossover_hardware import (
+ assert_hardware_report,
+ policy,
+ policy_for_report,
+)
+
+
+def main():
+ report = {
+ "hardware_profile": "a6000",
+ "cuda_visible_devices": "3",
+ "physical_gpu_index": 3,
+ "physical_gpu_uuid": "GPU-smoke",
+ "physical_gpu_name": "NVIDIA RTX A6000",
+ "total_memory_mib": 49140,
+ }
+ a6000 = policy_for_report("a6000", report)
+ assert a6000["allowed_physical_gpu_indices"] == [3]
+ assert_hardware_report(report, a6000)
+ gtx = policy("timan107_gtx1080")
+ assert gtx["allowed_physical_gpu_indices"] == [5, 7]
+ assert gtx["allowed_device_names"] == ["NVIDIA GeForce GTX 1080"]
+ print("crossover hardware smoke: all checks passed")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/experiments/prepare_shakespeare_char.py b/experiments/prepare_shakespeare_char.py
new file mode 100755
index 0000000..f09fc80
--- /dev/null
+++ b/experiments/prepare_shakespeare_char.py
@@ -0,0 +1,77 @@
+#!/usr/bin/env python3
+"""Download and deterministically tokenize the frozen Tiny Shakespeare data."""
+import argparse
+import hashlib
+import os
+import pickle
+import urllib.request
+
+import numpy as np
+
+
+URL = (
+ "https://raw.githubusercontent.com/karpathy/char-rnn/master/"
+ "data/tinyshakespeare/input.txt"
+)
+EXPECTED = {
+ "input.txt":
+ "86c4e6aa9db7c042ec79f339dcb96d42b0075e16b8fc2e86bf0ca57e2dc565ed",
+ "train.bin":
+ "6ec305602a99ac2802745a134e1f5e33e2231b4855525b00b9aebb730ac2626f",
+ "val.bin":
+ "d37d30cc0c8327c270d493299c3dca54135f6d5f1c9ef60cda78076e311204b1",
+}
+
+
+def sha256(path):
+ digest = hashlib.sha256()
+ with open(path, "rb") as handle:
+ for chunk in iter(lambda: handle.read(1024 * 1024), b""):
+ digest.update(chunk)
+ return digest.hexdigest()
+
+
+def verify(path, expected):
+ actual = sha256(path)
+ if actual != expected:
+ raise RuntimeError(
+ f"hash mismatch for {path}: expected {expected}, found {actual}"
+ )
+
+
+def main():
+ parser = argparse.ArgumentParser()
+ parser.add_argument("--out-dir", required=True)
+ args = parser.parse_args()
+ root = os.path.abspath(args.out_dir)
+ os.makedirs(root, exist_ok=True)
+ input_path = os.path.join(root, "input.txt")
+ if not os.path.isfile(input_path):
+ urllib.request.urlretrieve(URL, input_path)
+ verify(input_path, EXPECTED["input.txt"])
+
+ with open(input_path, encoding="utf-8") as handle:
+ data = handle.read()
+ chars = sorted(set(data))
+ stoi = {char: index for index, char in enumerate(chars)}
+ itos = {index: char for char, index in stoi.items()}
+ encoded = np.asarray([stoi[char] for char in data], dtype=np.uint16)
+ boundary = int(len(encoded) * 0.9)
+ encoded[:boundary].tofile(os.path.join(root, "train.bin"))
+ encoded[boundary:].tofile(os.path.join(root, "val.bin"))
+ with open(os.path.join(root, "meta.pkl"), "wb") as handle:
+ pickle.dump(
+ {"vocab_size": len(chars), "itos": itos, "stoi": stoi},
+ handle,
+ )
+
+ for name, expected in EXPECTED.items():
+ verify(os.path.join(root, name), expected)
+ print(
+ f"prepared Tiny Shakespeare at {root}: "
+ f"{boundary} train, {len(encoded) - boundary} validation tokens"
+ )
+
+
+if __name__ == "__main__":
+ main()
diff --git a/experiments/resnet_crossover_r2.py b/experiments/resnet_crossover_r2.py
index a9d5dfe..a7e78d7 100644
--- a/experiments/resnet_crossover_r2.py
+++ b/experiments/resnet_crossover_r2.py
@@ -10,6 +10,15 @@ import time
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+sys.path.insert(0, ROOT)
+from experiments.crossover_hardware import (
+ DEFAULT_PROFILE,
+ physical_gpu_report,
+ policy_for_report,
+ profile_choices,
+)
+
+
PROTOCOL = os.path.join(ROOT, "RESNET_CROSSOVER.md")
RESULT_ROOT = os.path.join(ROOT, "results", "resnet_crossover")
DEFAULT_SELECTOR = os.path.join(RESULT_ROOT, "p1_selector.json")
@@ -113,6 +122,7 @@ def source_report(selector_path):
os.path.join(ROOT, "experiments", "analyze_resnet_crossover_r2.py"),
os.path.join(ROOT, "experiments", "resnet_crossover_native.py"),
os.path.join(ROOT, "experiments", "conv_run.py"),
+ os.path.join(ROOT, "experiments", "crossover_hardware.py"),
os.path.join(ROOT, "sdil", "conv_crossover.py"),
os.path.join(ROOT, "sdil", "conv.py"),
os.path.join(ROOT, "sdil", "data.py"),
@@ -295,7 +305,7 @@ def registry_sha256(jobs):
return hashlib.sha256(encoded).hexdigest()
-def ensure_launch(source, selector, jobs):
+def ensure_launch(source, selector, jobs, hardware_policy):
path = os.path.join(RESULT_ROOT, "r2_launch.json")
expected = {
"stage": "r2",
@@ -303,7 +313,9 @@ def ensure_launch(source, selector, jobs):
"selector": selector,
"registry_sha256": registry_sha256(jobs),
"num_jobs": len(jobs),
- "allowed_physical_gpus": [5, 7],
+ "hardware_policy": hardware_policy,
+ "allowed_physical_gpus":
+ hardware_policy["allowed_physical_gpu_indices"],
}
if os.path.isfile(path):
if read_json(path) != expected:
@@ -326,40 +338,6 @@ def assert_source_unchanged(source):
raise RuntimeError(f"ResNet R2 source drift: {relative}")
-def physical_gpu_report(dry_run=False):
- visible = os.environ.get("CUDA_VISIBLE_DEVICES")
- if dry_run:
- return {
- "cuda_visible_devices": visible,
- "physical_gpu_index": None,
- "physical_gpu_uuid": None,
- }
- if visible not in {"5", "7"}:
- raise RuntimeError("ResNet R2 requires physical GPU 5 or 7")
- query = subprocess.run(
- [
- "nvidia-smi",
- "--query-gpu=index,uuid,name",
- "--format=csv,noheader,nounits",
- ],
- check=True,
- capture_output=True,
- text=True,
- ).stdout.splitlines()
- rows = {}
- for line in query:
- index, uuid, name = [value.strip() for value in line.split(",", 2)]
- rows[index] = {"uuid": uuid, "name": name}
- if visible not in rows:
- raise RuntimeError(f"physical GPU {visible} not found")
- return {
- "cuda_visible_devices": visible,
- "physical_gpu_index": int(visible),
- "physical_gpu_uuid": rows[visible]["uuid"],
- "physical_gpu_name": rows[visible]["name"],
- }
-
-
def run_job(job, source, selector, gpu, dry_run):
if not dry_run:
assert_source_unchanged(source)
@@ -414,16 +392,24 @@ def main():
parser.add_argument("--num-shards", type=int, default=1)
parser.add_argument("--method", choices=METHODS)
parser.add_argument("--depth", type=int, choices=DEPTHS)
+ parser.add_argument(
+ "--hardware-profile",
+ choices=profile_choices(),
+ default=os.environ.get("SDIL_HARDWARE_PROFILE", DEFAULT_PROFILE),
+ )
parser.add_argument("--dry-run", action="store_true")
args = parser.parse_args()
if not 0 <= args.shard_index < args.num_shards:
raise ValueError("invalid ResNet R2 shard")
selector = selector_report(args.selector)
source = source_report(args.selector)
- gpu = physical_gpu_report(args.dry_run)
+ gpu = physical_gpu_report(args.hardware_profile, args.dry_run)
complete_jobs = r2_jobs(selector["selected_rates"])
if not args.dry_run:
- launch = ensure_launch(source, selector, complete_jobs)
+ hardware_policy = policy_for_report(args.hardware_profile, gpu)
+ launch = ensure_launch(
+ source, selector, complete_jobs, hardware_policy
+ )
print(f"R2 launch lock: {launch}", flush=True)
jobs = complete_jobs
if args.method:
diff --git a/experiments/run_a6000_matrix.py b/experiments/run_a6000_matrix.py
new file mode 100755
index 0000000..e5330ef
--- /dev/null
+++ b/experiments/run_a6000_matrix.py
@@ -0,0 +1,182 @@
+#!/usr/bin/env python3
+"""Sequential, restart-safe launcher for the complete A6000 81-cell matrix."""
+import argparse
+import os
+import subprocess
+import sys
+
+
+ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+sys.path.insert(0, ROOT)
+from experiments.crossover_hardware import physical_gpu_report
+from experiments.resnet_crossover_r2 import selector_report as resnet_selector
+from experiments.transformer_crossover_t2 import (
+ selector_report as transformer_selector,
+)
+
+
+def run(command, cwd, environment, dry_run):
+ rendered = " ".join(command)
+ print(f"[{cwd}] {rendered}", flush=True)
+ if not dry_run:
+ subprocess.run(command, cwd=cwd, env=environment, check=True)
+
+
+def require_file(path, label):
+ if not os.path.isfile(path):
+ raise RuntimeError(f"{label} is not frozen in this run packet: {path}")
+
+
+def main():
+ parser = argparse.ArgumentParser()
+ parser.add_argument("--plain-root", required=True)
+ parser.add_argument("--plain-python", required=True)
+ parser.add_argument(
+ "--family",
+ choices=("all", "plain", "resnet", "transformer"),
+ default="all",
+ )
+ parser.add_argument("--dry-run", action="store_true")
+ args = parser.parse_args()
+
+ plain_root = os.path.abspath(args.plain_root)
+ plain_python = os.path.abspath(args.plain_python)
+ require_file(plain_python, "plain-CNN Python executable")
+ require_file(
+ os.path.join(plain_root, "crossover_grid.py"),
+ "patched plain-CNN runner",
+ )
+ gpu = physical_gpu_report("a6000", args.dry_run)
+ print(f"hardware lock: {gpu}", flush=True)
+
+ environment = os.environ.copy()
+ environment["SDIL_HARDWARE_PROFILE"] = "a6000"
+ environment["SDIL_MAIN_ROOT"] = ROOT
+ families = (
+ ("plain", "resnet", "transformer")
+ if args.family == "all" else (args.family,)
+ )
+
+ plain_audit = os.path.join(plain_root, "runs", "plain-p2-audit.json")
+ if "plain" in families:
+ run(
+ [
+ plain_python,
+ "tests/local_rules_smoke.py",
+ ],
+ plain_root,
+ environment,
+ args.dry_run,
+ )
+ run(
+ [
+ plain_python,
+ "crossover_grid.py",
+ "--stage", "p2",
+ "--hardware-profile", "a6000",
+ ],
+ plain_root,
+ environment,
+ args.dry_run,
+ )
+ run(
+ [
+ plain_python,
+ "analyze_p2.py",
+ "--output", plain_audit,
+ ],
+ plain_root,
+ environment,
+ args.dry_run,
+ )
+
+ resnet_selector_path = os.path.join(
+ ROOT, "results", "resnet_crossover", "p1_selector.json"
+ )
+ if "resnet" in families:
+ require_file(resnet_selector_path, "complete ResNet P1 selector")
+ if not args.dry_run:
+ resnet_selector(resnet_selector_path)
+ run(
+ [
+ sys.executable,
+ "experiments/resnet_crossover_smoke.py",
+ ],
+ ROOT,
+ environment,
+ args.dry_run,
+ )
+ run(
+ [
+ sys.executable,
+ "experiments/resnet_crossover_r2.py",
+ "--hardware-profile", "a6000",
+ ],
+ ROOT,
+ environment,
+ args.dry_run,
+ )
+ run(
+ [
+ sys.executable,
+ "experiments/analyze_resnet_crossover_r2.py",
+ ],
+ ROOT,
+ environment,
+ args.dry_run,
+ )
+
+ transformer_selector_path = os.path.join(
+ ROOT, "results", "transformer_crossover", "p1_selector.json"
+ )
+ if "transformer" in families:
+ require_file(
+ transformer_selector_path, "complete Transformer T1 selector"
+ )
+ if not args.dry_run:
+ transformer_selector(transformer_selector_path)
+ run(
+ [
+ sys.executable,
+ "experiments/transformer_feedback_smoke.py",
+ ],
+ ROOT,
+ environment,
+ args.dry_run,
+ )
+ run(
+ [
+ sys.executable,
+ "experiments/transformer_crossover_t2.py",
+ "--hardware-profile", "a6000",
+ ],
+ ROOT,
+ environment,
+ args.dry_run,
+ )
+ run(
+ [
+ sys.executable,
+ "experiments/analyze_transformer_crossover_t2.py",
+ ],
+ ROOT,
+ environment,
+ args.dry_run,
+ )
+
+ if args.family == "all":
+ require_file(plain_audit, "complete plain-CNN P2 audit")
+ run(
+ [
+ sys.executable,
+ "experiments/analyze_crossover_81.py",
+ "--plain", plain_audit,
+ ],
+ ROOT,
+ environment,
+ args.dry_run,
+ )
+
+
+if __name__ == "__main__":
+ main()
diff --git a/experiments/transformer_crossover_native.py b/experiments/transformer_crossover_native.py
index ca47751..6a0269e 100644
--- a/experiments/transformer_crossover_native.py
+++ b/experiments/transformer_crossover_native.py
@@ -25,7 +25,10 @@ METHODS = (
"bp", "fa", "dfa", "pepita", "ff", "ep", "dualprop",
"clean_kp", "sdil",
)
-DEFAULT_DATA_DIR = "/home/yurenh2/local-transformer/data/shakespeare_char"
+DEFAULT_DATA_DIR = os.environ.get(
+ "SDIL_TRANSFORMER_DATA_DIR",
+ "/home/yurenh2/local-transformer/data/shakespeare_char",
+)
def file_sha256(path):
diff --git a/experiments/transformer_crossover_t2.py b/experiments/transformer_crossover_t2.py
index 17f61a2..4f12e56 100644
--- a/experiments/transformer_crossover_t2.py
+++ b/experiments/transformer_crossover_t2.py
@@ -10,6 +10,15 @@ import time
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+sys.path.insert(0, ROOT)
+from experiments.crossover_hardware import (
+ DEFAULT_PROFILE,
+ physical_gpu_report,
+ policy_for_report,
+ profile_choices,
+)
+
+
PROTOCOL = os.path.join(ROOT, "TRANSFORMER_CROSSOVER.md")
RESULT_ROOT = os.path.join(ROOT, "results", "transformer_crossover")
DEFAULT_SELECTOR = os.path.join(RESULT_ROOT, "p1_selector.json")
@@ -123,6 +132,7 @@ def source_report(selector_path):
os.path.join(
ROOT, "experiments", "transformer_feedback_smoke.py"
),
+ os.path.join(ROOT, "experiments", "crossover_hardware.py"),
os.path.join(ROOT, "sdil", "transformer.py"),
os.path.abspath(selector_path),
]
@@ -235,7 +245,7 @@ def registry_sha256(jobs):
return hashlib.sha256(encoded).hexdigest()
-def ensure_launch(source, selector, jobs):
+def ensure_launch(source, selector, jobs, hardware_policy):
path = os.path.join(RESULT_ROOT, "t2_launch.json")
expected = {
"stage": "t2",
@@ -243,7 +253,9 @@ def ensure_launch(source, selector, jobs):
"selector": selector,
"registry_sha256": registry_sha256(jobs),
"num_jobs": len(jobs),
- "allowed_physical_gpus": [5, 7],
+ "hardware_policy": hardware_policy,
+ "allowed_physical_gpus":
+ hardware_policy["allowed_physical_gpu_indices"],
}
if os.path.isfile(path):
if read_json(path) != expected:
@@ -274,40 +286,6 @@ def assert_source_unchanged(source):
raise RuntimeError("Transformer T2 environment drift")
-def physical_gpu_report(dry_run=False):
- visible = os.environ.get("CUDA_VISIBLE_DEVICES")
- if dry_run:
- return {
- "cuda_visible_devices": visible,
- "physical_gpu_index": None,
- "physical_gpu_uuid": None,
- }
- if visible not in {"5", "7"}:
- raise RuntimeError("Transformer T2 requires physical GPU 5 or 7")
- query = subprocess.run(
- [
- "nvidia-smi",
- "--query-gpu=index,uuid,name",
- "--format=csv,noheader,nounits",
- ],
- check=True,
- capture_output=True,
- text=True,
- ).stdout.splitlines()
- rows = {}
- for line in query:
- index, uuid, name = [value.strip() for value in line.split(",", 2)]
- rows[index] = {"uuid": uuid, "name": name}
- if visible not in rows:
- raise RuntimeError(f"physical GPU {visible} not found")
- return {
- "cuda_visible_devices": visible,
- "physical_gpu_index": int(visible),
- "physical_gpu_uuid": rows[visible]["uuid"],
- "physical_gpu_name": rows[visible]["name"],
- }
-
-
def run_job(job, source, selector, gpu, dry_run):
if not dry_run:
assert_source_unchanged(source)
@@ -361,16 +339,24 @@ def main():
parser.add_argument("--num-shards", type=int, default=1)
parser.add_argument("--method", choices=METHODS)
parser.add_argument("--depth", type=int, choices=DEPTHS)
+ parser.add_argument(
+ "--hardware-profile",
+ choices=profile_choices(),
+ default=os.environ.get("SDIL_HARDWARE_PROFILE", DEFAULT_PROFILE),
+ )
parser.add_argument("--dry-run", action="store_true")
args = parser.parse_args()
if not 0 <= args.shard_index < args.num_shards:
raise ValueError("invalid Transformer T2 shard")
selector = selector_report(args.selector)
source = source_report(args.selector)
- gpu = physical_gpu_report(args.dry_run)
+ gpu = physical_gpu_report(args.hardware_profile, args.dry_run)
complete_jobs = t2_jobs(selector["selected_rates"])
if not args.dry_run:
- launch = ensure_launch(source, selector, complete_jobs)
+ hardware_policy = policy_for_report(args.hardware_profile, gpu)
+ launch = ensure_launch(
+ source, selector, complete_jobs, hardware_policy
+ )
print(f"T2 launch lock: {launch}", flush=True)
jobs = complete_jobs
if args.method:
diff --git a/external/README.md b/external/README.md
new file mode 100644
index 0000000..799a63b
--- /dev/null
+++ b/external/README.md
@@ -0,0 +1,11 @@
+# External baseline reconstruction
+
+`dualprop_patches/` is an ordered `git format-patch` series on top of the
+Dual Propagation authors' public revision
+`7b2595b34421e1483a721dbfdeff8cdabda3a1ff`. It contains the matched
+plain-CNN adapters and the portable A6000 hardware profile; it does not vendor
+the authors' repository.
+
+Run `experiments/bootstrap_plain_cnn.sh /absolute/target/path` to clone the
+author repository, check out the frozen revision, and apply all 19 patches in
+order. The resulting repository must remain clean during formal runs.
diff --git a/external/dualprop_patches/0001-crossover-separate-author-diagnostics-from-timed-tra.patch b/external/dualprop_patches/0001-crossover-separate-author-diagnostics-from-timed-tra.patch
new file mode 100644
index 0000000..6340b61
--- /dev/null
+++ b/external/dualprop_patches/0001-crossover-separate-author-diagnostics-from-timed-tra.patch
@@ -0,0 +1,198 @@
+From cd932728011d5fef91109dd24df8a5f2dfb9e6e7 Mon Sep 17 00:00:00 2001
+From: YurenHao0426 <Blackhao0426@gmail.com>
+Date: Mon, 27 Jul 2026 12:50:47 -0500
+Subject: [PATCH 01/19] crossover: separate author diagnostics from timed
+ training
+
+---
+ config/cli_config.py | 14 +++++++++++++-
+ src/training_utils.py | 42 ++++++++++++++++++++++++++++++------------
+ train.py | 27 ++++++++++++++++++++-------
+ 3 files changed, 63 insertions(+), 20 deletions(-)
+
+diff --git a/config/cli_config.py b/config/cli_config.py
+index 0aec5f6..c1b23ba 100644
+--- a/config/cli_config.py
++++ b/config/cli_config.py
+@@ -44,6 +44,18 @@ parser.add_argument('--model', default='VGG16', choices=['VGG16', 'VGGlike', 'CN
+
+ parser.add_argument('--learning-algorithm', default='dualprop-lagr-ff', choices=['backprop', 'dualprop-lagr-ff', 'dualprop-raovr-ff', 'dualprop-raovr-dampened-ff'])
+
++parser.add_argument(
++ '--gradient-diagnostics', default='full', choices=['none', 'full'],
++ help=('Compute the exact BP reference gradient and layerwise cosine on '
++ 'every training minibatch. The author-compatible default is full; '
++ 'use none for diagnostic-free timed crossover runs.'))
++
++parser.add_argument(
++ '--spectral-diagnostics', default='full', choices=['none', 'full'],
++ help=('Run the author power-iteration L/gamma probes after every '
++ 'validation evaluation. The author-compatible default is full; '
++ 'use none for diagnostic-free timed crossover runs.'))
++
+ dtypes = {'bfloat16': jnp.bfloat16, 'float16':jnp.float16, 'float32':jnp.float32}
+ parser.add_argument('--dtype', default='float32', choices=dtypes.keys())
+ parser.add_argument('--param-dtype', default='float32', choices=['bfloat16', 'float16', 'float32'])
+@@ -120,4 +132,4 @@ config.model = modeltype[config.learning_algorithm](loss_func, Conv, Dense, acti
+ )
+
+ # Load datasets
+-config.train_ds, config.val_ds, config.test_ds = datasets[config.dataset](config.dtype, config.percent_train, config.percent_val)
+\ No newline at end of file
++config.train_ds, config.val_ds, config.test_ds = datasets[config.dataset](config.dtype, config.percent_train, config.percent_val)
+diff --git a/src/training_utils.py b/src/training_utils.py
+index 7152451..2fbfccd 100644
+--- a/src/training_utils.py
++++ b/src/training_utils.py
+@@ -240,7 +240,8 @@ def to_float16(ptree):
+ def to_float32(ptree):
+ return tree_map(lambda x: x.astype(jnp.float32), ptree)
+
+-def train_epoch(state, train_ds, batch_size, rng, augmentation_on, learning_algorithm, num_classes):
++def train_epoch(state, train_ds, batch_size, rng, augmentation_on,
++ learning_algorithm, num_classes, gradient_diagnostics=True):
+ """Train for a single epoch."""
+ t0 = time.time()
+ train_ds_size = len(train_ds['image'])
+@@ -261,7 +262,9 @@ def train_epoch(state, train_ds, batch_size, rng, augmentation_on, learning_algo
+ # image = vmap_augment_train_imagenet(image, batch_rng)
+
+ if learning_algorithm != "backprop":
+- state, metrics = train_step(state, image, labels_onehot, labels, batch_rng, inf_rng, augmentation_on)
++ state, metrics = train_step(
++ state, image, labels_onehot, labels, batch_rng, inf_rng,
++ augmentation_on, gradient_diagnostics)
+ elif learning_algorithm == "backprop":
+ state, metrics = train_step_bp(state, image, labels_onehot, labels, batch_rng, inf_rng, augmentation_on)
+ batch_metrics.append(metrics)
+@@ -279,7 +282,8 @@ def train_epoch(state, train_ds, batch_size, rng, augmentation_on, learning_algo
+ return state, epoch_metrics_np, runtime
+
+ @jax.jit
+-def train_step(state, image, labels_onehot, labels, batch_rng, inf_rng, augmentation_on):
++def train_step(state, image, labels_onehot, labels, batch_rng, inf_rng,
++ augmentation_on, gradient_diagnostics):
+ """Train for a single step."""
+ # batch_rng = jax.random.split(batch_rng, batch['image'].shape[0])
+
+@@ -298,8 +302,9 @@ def train_step(state, image, labels_onehot, labels, batch_rng, inf_rng, augmenta
+ inf_rng, _ = jax.random.split(inf_rng)
+ metrics = compute_metrics(image=image, labels_onehot=labels_onehot, labels=labels, state=state)
+
+- get_ref_grad_angle = True
+- metrics = jax.lax.cond(get_ref_grad_angle, ref_grad_and_angle, no_ref_grad_and_angle, state, grads, image, labels_onehot, metrics)
++ metrics = jax.lax.cond(
++ gradient_diagnostics, ref_grad_and_angle, no_ref_grad_and_angle,
++ state, grads, image, labels_onehot, metrics)
+
+ # The optimizer may modify grads, so we need to compare grads and ref_grads before performing the gradient step.
+ state = state.apply_gradients(grads=grads)
+@@ -340,7 +345,8 @@ def eval_step(state, params, image, labels_onehot, labels, inf_rng):
+ metrics = compute_metrics(image=image, labels_onehot=labels_onehot, labels=labels, state=state)
+ return metrics
+
+-def eval_model(state, params, test_ds, batch_size, num_classes, eval_rng):
++def eval_model(state, params, test_ds, batch_size, num_classes, eval_rng,
++ spectral_diagnostics=True):
+ t0 = time.time()
+ test_ds_size = len(test_ds['image'])
+ steps = test_ds_size // batch_size
+@@ -365,11 +371,23 @@ def eval_model(state, params, test_ds, batch_size, num_classes, eval_rng):
+
+ runtime = time.time() - t0
+
+- # dummy states, used by get_L_and_gamma to infer correct array shape when generating random arrays
+- sdummy = state.apply_fn({'params': params}, batch["image"], method='make_predictions')
+- eval_rng, _ = jax.random.split(eval_rng)
+- L10, gamma10 = state.apply_fn({'params': state.params}, s=sdummy, rng_key=eval_rng, numiter=10, method='get_L_and_gamma')
+- L20, gamma20 = state.apply_fn({'params': state.params}, s=sdummy, rng_key=eval_rng, numiter=20, method='get_L_and_gamma')
++ if spectral_diagnostics:
++ # Dummy states infer the array shapes used by the author's power
++ # iterations. These diagnostics are expensive and are deliberately
++ # optional in timed crossover runs.
++ sdummy = state.apply_fn(
++ {'params': params}, batch["image"], method='make_predictions')
++ eval_rng, _ = jax.random.split(eval_rng)
++ L10, gamma10 = state.apply_fn(
++ {'params': state.params}, s=sdummy, rng_key=eval_rng, numiter=10,
++ method='get_L_and_gamma')
++ L20, gamma20 = state.apply_fn(
++ {'params': state.params}, s=sdummy, rng_key=eval_rng, numiter=20,
++ method='get_L_and_gamma')
++ else:
++ count = len(state.params)
++ L10 = L20 = gamma10 = gamma20 = [
++ jnp.asarray(jnp.nan, dtype=jnp.float32) for _ in range(count)]
+
+ return summary['loss'], summary['accuracy'], summary['top5accuracy'], runtime, L10, L20, gamma10, gamma20
+
+@@ -470,4 +488,4 @@ def plot_L_or_gamma(L20, L10, ylabel, save_path):
+ plt.tight_layout()
+ plt.savefig(save_path)
+ plt.close()
+- return
+\ No newline at end of file
++ return
+diff --git a/train.py b/train.py
+index bd7f6b7..139fd6c 100644
+--- a/train.py
++++ b/train.py
+@@ -71,7 +71,10 @@ for experiment_index, seed in enumerate(config.seeds):
+ # Run an optimization step over a training batch
+ # last augument turns off data augmentation for mnist
+ augmentation_on = (config.dataset!="mnist") and (config.dataset!="fashionmnist")
+- state, epoch_metrics, train_time = train_epoch(state, config.train_ds, config.batch_size, input_rng, augmentation_on, config.learning_algorithm, config.num_classes)
++ state, epoch_metrics, train_time = train_epoch(
++ state, config.train_ds, config.batch_size, input_rng,
++ augmentation_on, config.learning_algorithm, config.num_classes,
++ gradient_diagnostics=(config.gradient_diagnostics == "full"))
+ loginfo_and_print('train: \tloss: %.4f, \taccuracy: %.4f, \truntime: %.4f' % (epoch_metrics["loss"], epoch_metrics["accuracy"], train_time))
+ hist['train_loss'][epoch-1], hist['train_accuracy'][epoch-1], hist['train_time'][epoch-1] = epoch_metrics["loss"], epoch_metrics["accuracy"], train_time
+
+@@ -83,7 +86,10 @@ for experiment_index, seed in enumerate(config.seeds):
+
+ # Evaluate on the validation set after each training epoch
+ rng, input_rng = jax.random.split(rng)
+- val_loss, val_accuracy, val_top5_accuracy, val_time, L10, L20, gamma10, gamma20 = eval_model(state, state.params, config.val_ds, config.batch_size, config.num_classes, input_rng)
++ val_loss, val_accuracy, val_top5_accuracy, val_time, L10, L20, gamma10, gamma20 = eval_model(
++ state, state.params, config.val_ds, config.batch_size,
++ config.num_classes, input_rng,
++ spectral_diagnostics=(config.spectral_diagnostics == "full"))
+ loginfo_and_print('val: \tloss: %.4f, \taccuracy: %.4f, \ttop5_accuracy: %.4f, \truntime: %.4f' % (val_loss, val_accuracy, val_top5_accuracy, val_time))
+ loginfo_and_print(f"L20: {[np.round(Li.item(), decimals=4) for Li in L20]}")
+ loginfo_and_print(f"gamma20: {[np.round(gi.item(), decimals=4) for gi in gamma20]}")
+@@ -108,18 +114,25 @@ for experiment_index, seed in enumerate(config.seeds):
+ loginfo_and_print(f"\n====Loading model with best validation accuracy (epoch {best_epoch})====")
+ best_state = checkpoints.restore_checkpoint(ckpt_dir=CKPT_DIR, target=state)
+ rng, input_rng = jax.random.split(rng)
+- test_loss, test_accuracy, test_top5accuracy, test_time, _, _, _, _ = eval_model(best_state, best_state.params, config.test_ds, config.batch_size, config.num_classes, input_rng)
++ test_loss, test_accuracy, test_top5accuracy, test_time, _, _, _, _ = eval_model(
++ best_state, best_state.params, config.test_ds, config.batch_size,
++ config.num_classes, input_rng,
++ spectral_diagnostics=(config.spectral_diagnostics == "full"))
+ hist['test_loss'], hist['test_accuracy'], hist['test_top5accuracy'], hist['test_time'] = test_loss, test_accuracy, test_top5accuracy, test_time
+ loginfo_and_print('test: \tloss: %.4f, \taccuracy: %.4f, \ttop5_accuracy: %.4f, \truntime: %.4f' % (test_loss, test_accuracy, test_top5accuracy, test_time))
+
+- if config.learning_algorithm != "backprop":
++ if (config.learning_algorithm != "backprop"
++ and config.gradient_diagnostics == "full"):
+ # Use color_norm=LogNorm(clip=True) for logscale plot
+ heatmap_grads_epochs(hist["grad_cos_sim_epochs"], outpath+"grad_angle_epochs.pdf", True, color_norm=None)
+ #Grad angle across batches in the first epoch
+ first_N = 100
+ heatmap_grads_batches(hist["grad_cos_sim_batches"][:,0:first_N], outpath+f"grad_angle_first_{first_N}_batches.pdf", True, color_norm=None)
+
+- plot_L_or_gamma(hist["L20"], hist["L10"], "L", outpath+"L.pdf")
+- plot_L_or_gamma(hist["gamma20"], hist["gamma10"], r"$\gamma$", outpath+"gamma.pdf")
++ if config.spectral_diagnostics == "full":
++ plot_L_or_gamma(hist["L20"], hist["L10"], "L", outpath+"L.pdf")
++ plot_L_or_gamma(
++ hist["gamma20"], hist["gamma10"], r"$\gamma$",
++ outpath+"gamma.pdf")
+
+- np.save(outpath+"hist.npy", hist)
+\ No newline at end of file
++ np.save(outpath+"hist.npy", hist)
+--
+2.54.0
+
diff --git a/external/dualprop_patches/0002-crossover-add-audited-ordinary-FA-and-DFA-rules.patch b/external/dualprop_patches/0002-crossover-add-audited-ordinary-FA-and-DFA-rules.patch
new file mode 100644
index 0000000..3138039
--- /dev/null
+++ b/external/dualprop_patches/0002-crossover-add-audited-ordinary-FA-and-DFA-rules.patch
@@ -0,0 +1,453 @@
+From ff5d9524a56d8b440c8524e6e571cf1488ce6aaa Mon Sep 17 00:00:00 2001
+From: YurenHao0426 <Blackhao0426@gmail.com>
+Date: Mon, 27 Jul 2026 13:03:33 -0500
+Subject: [PATCH 02/19] crossover: add audited ordinary FA and DFA rules
+
+---
+ config/cli_config.py | 8 ++-
+ src/__init__.py | 2 +-
+ src/models.py | 80 ++++++++++++++++++++-
+ src/training_utils.py | 90 +++++++++++++++++++++++-
+ tests/local_rules_smoke.py | 138 +++++++++++++++++++++++++++++++++++++
+ train.py | 6 +-
+ 6 files changed, 317 insertions(+), 7 deletions(-)
+ create mode 100644 tests/local_rules_smoke.py
+
+diff --git a/config/cli_config.py b/config/cli_config.py
+index c1b23ba..e1756e1 100644
+--- a/config/cli_config.py
++++ b/config/cli_config.py
+@@ -42,7 +42,11 @@ parser.add_argument('--experiment-name', default='test', help='A string denoting
+
+ parser.add_argument('--model', default='VGG16', choices=['VGG16', 'VGGlike', 'CNN', 'miniCNN', 'MLP'], help='')
+
+-parser.add_argument('--learning-algorithm', default='dualprop-lagr-ff', choices=['backprop', 'dualprop-lagr-ff', 'dualprop-raovr-ff', 'dualprop-raovr-dampened-ff'])
++parser.add_argument('--learning-algorithm', default='dualprop-lagr-ff', choices=['backprop', 'fa', 'dfa', 'dualprop-lagr-ff', 'dualprop-raovr-ff', 'dualprop-raovr-dampened-ff'])
++
++parser.add_argument(
++ '--feedback-seed', default=1729, type=int,
++ help='Independent fixed-feedback initialization seed for FA/DFA.')
+
+ parser.add_argument(
+ '--gradient-diagnostics', default='full', choices=['none', 'full'],
+@@ -124,7 +128,7 @@ elif config.model == "MLP":
+ dense_features = [1024, 1024, config.num_classes]
+
+ # Load model
+-modeltype = {"backprop":cnn_abstract, "dualprop-lagr-ff": cnn_dualprop_Lagr_ff, "dualprop-raovr-ff": cnn_dualprop_RAOVR_ff, "dualprop-raovr-dampened-ff": cnn_dualprop_RAOVR_dampened_ff}
++modeltype = {"backprop":cnn_abstract, "fa": cnn_abstract, "dfa": cnn_abstract, "dualprop-lagr-ff": cnn_dualprop_Lagr_ff, "dualprop-raovr-ff": cnn_dualprop_RAOVR_ff, "dualprop-raovr-dampened-ff": cnn_dualprop_RAOVR_dampened_ff}
+ activation={"relu": relu, "hs": hs, "sigmoid": sigmoid, "tanh": tanh}
+ config.model = modeltype[config.learning_algorithm](loss_func, Conv, Dense, activation[config.activation], config.num_classes, config.beta, config.alpha, config.dtype, config.param_dtype,
+ kernels=kernels, strides=strides, features=features, mp = mp,
+diff --git a/src/__init__.py b/src/__init__.py
+index f8ce731..7a08c87 100644
+--- a/src/__init__.py
++++ b/src/__init__.py
+@@ -1,2 +1,2 @@
+ from .models import cnn_dualprop_Lagr_ff, cnn_dualprop_RAOVR_ff, cnn_dualprop_RAOVR_dampened_ff, cnn_abstract
+-from .training_utils import create_train_state, train_epoch, eval_model, get_mnist, get_svhn, get_fashionmnist, get_cifar10, get_cifar100, get_imagenet_32x32, heatmap_grads_batches, heatmap_grads_epochs, plot_L_or_gamma
+\ No newline at end of file
++from .training_utils import create_train_state, create_local_feedback, train_epoch, eval_model, get_mnist, get_svhn, get_fashionmnist, get_cifar10, get_cifar100, get_imagenet_32x32, heatmap_grads_batches, heatmap_grads_epochs, plot_L_or_gamma
+diff --git a/src/models.py b/src/models.py
+index c162209..edc432c 100644
+--- a/src/models.py
++++ b/src/models.py
+@@ -116,6 +116,84 @@ class cnn_abstract(nn.Module, ABC):
+ s.append(self.act(layer(s[-1])))
+ s.append(self.layers[-1](s[-1]))
+ return s
++
++ def ff_with_local_cache(self, x0):
++ """Forward pass with the linear outputs needed by audited local rules.
++
++ ``s[i]`` is the post-nonlinearity input to layer ``i`` and
++ ``linear[i]`` is that layer's output before its parameter-free
++ pooling/nonlinearity. Keeping the two pieces separate lets feedback
++ alignment use the *forward* ReLU mask and max-pool switches while
++ replacing only the learned linear transpose by an independent random
++ operator.
++ """
++ s = [x0]
++ linear = []
++ for i, layer in enumerate(self.layers):
++ if i < self.num_convlayers:
++ z = layer.call_without_pooling(s[-1])
++ h = self.act(layer.pooling(z))
++ else:
++ z = layer(s[-1])
++ h = z if i == self.num_layers - 1 else self.act(z)
++ linear.append(z)
++ s.append(h)
++ return s, linear
++
++ def local_correlation_objective(self, s, teaching_fields):
++ """Sum independent per-layer correlations for a local weight update.
++
++ Both the layer inputs and teaching fields are detached. Differentiating
++ this scalar with respect to the model parameters therefore evaluates
++ only each layer's eligibility/Jacobian and cannot create a reverse
++ graph through another forward layer.
++ """
++ objective = 0.0
++ batch_size = s[0].shape[0]
++ for i, layer in enumerate(self.layers):
++ x = jax.lax.stop_gradient(s[i])
++ field = jax.lax.stop_gradient(teaching_fields[i + 1])
++ if i < self.num_convlayers:
++ z = layer.call_without_pooling(x)
++ h = self.act(layer.pooling(z))
++ else:
++ z = layer(x)
++ h = z if i == self.num_layers - 1 else self.act(z)
++ objective += jnp.sum(h * field)
++ return objective / batch_size
++
++ def fa_teaching_fields(self, s, linear, output_field):
++ """Transport a post-output field through fixed random linear maps.
++
++ This method is evaluated with an independent parameter tree. The
++ activation and max-pool pullbacks are evaluated at the cached forward
++ linear outputs, while only the convolution/dense transpose comes from
++ this method's parameters. In the audit-only symmetric limit where the
++ parameter tree equals the forward tree, the resulting fields are exact
++ reverse-mode fields.
++ """
++ fields = [jnp.zeros_like(value) for value in s]
++ fields[-1] = output_field
++ for i in range(self.num_layers - 1, 0, -1):
++ child_field = fields[i + 1]
++ if i == self.num_layers - 1:
++ linear_field = child_field
++ elif i < self.num_convlayers:
++ _, post_pullback = jax.vjp(
++ lambda z: self.act(self.layers[i].pooling(z)),
++ linear[i])
++ linear_field = post_pullback(child_field)[0]
++ else:
++ _, post_pullback = jax.vjp(self.act, linear[i])
++ linear_field = post_pullback(child_field)[0]
++
++ if i < self.num_convlayers:
++ linear_map = lambda x: self.layers[i].call_without_pooling(x)
++ else:
++ linear_map = self.layers[i]
++ _, linear_pullback = jax.vjp(linear_map, s[i])
++ fields[i] = linear_pullback(linear_field)[0]
++ return fields
+
+ def init_states_to_zero(self, x0):
+ s = [x0]
+@@ -327,4 +405,4 @@ class cnn_dualprop_RAOVR_dampened_ff(cnn_dualprop_abstract):
+ # delta = s[i+1] - fa[i+1]
+ # s[i], fa[i] = self.infer_hidden(s[i], s[i-1], delta, self.layers[i-1], self.layers[i], L[i])
+
+-# return s, fa
+\ No newline at end of file
++# return s, fa
+diff --git a/src/training_utils.py b/src/training_utils.py
+index 2fbfccd..b6e1e3c 100644
+--- a/src/training_utils.py
++++ b/src/training_utils.py
+@@ -15,6 +15,7 @@ import time, datetime # measuring runetime and generating timestamps for experim
+ import os # for os.makdirs() function
+ import seaborn as sns
+ import matplotlib.pylab as plt
++from functools import partial
+
+ class SumGreaterThan100Error(Exception):
+ pass
+@@ -216,6 +217,33 @@ def create_train_state(rng, model, image_dims, lr, wlr, lrf, momentum, weight_de
+
+ return train_state.TrainState.create(apply_fn=model.apply, params=unfreeze(params), tx=tx)
+
++
++def create_local_feedback(rng, model, params, image_dims, learning_algorithm,
++ num_classes):
++ """Create fixed feedback without reading a forward parameter value.
++
++ FA uses an independently initialized model-shaped linear parameter tree.
++ DFA uses one independent output-to-hidden tensor per hidden population.
++ The forward tree is used only to infer activation shapes for DFA.
++ """
++ if learning_algorithm not in ("fa", "dfa"):
++ return None
++ w, h, channels = image_dims
++ dummy = jnp.ones([1, w, h, channels])
++ if learning_algorithm == "fa":
++ return unfreeze(model.init(rng, dummy)["params"])
++
++ states, _ = model.apply(
++ {"params": params}, dummy, method="ff_with_local_cache")
++ keys = jax.random.split(rng, len(states) - 2)
++ scale = jnp.asarray(num_classes ** -0.5, dtype=dummy.dtype)
++ return tuple(
++ scale * jax.random.normal(
++ key, (num_classes,) + tuple(state.shape[1:]),
++ dtype=state.dtype)
++ for key, state in zip(keys, states[1:-1])
++ )
++
+ def augment_train(image, batch_rng):
+ w, h, c = image.shape
+
+@@ -241,7 +269,8 @@ def to_float32(ptree):
+ return tree_map(lambda x: x.astype(jnp.float32), ptree)
+
+ def train_epoch(state, train_ds, batch_size, rng, augmentation_on,
+- learning_algorithm, num_classes, gradient_diagnostics=True):
++ learning_algorithm, num_classes, local_feedback=None,
++ gradient_diagnostics=True):
+ """Train for a single epoch."""
+ t0 = time.time()
+ train_ds_size = len(train_ds['image'])
+@@ -261,7 +290,12 @@ def train_epoch(state, train_ds, batch_size, rng, augmentation_on,
+ batch_rng = jax.random.split(batch_rng, image.shape[0])
+ # image = vmap_augment_train_imagenet(image, batch_rng)
+
+- if learning_algorithm != "backprop":
++ if learning_algorithm in ("fa", "dfa"):
++ state, metrics = train_step_local(
++ state, local_feedback, learning_algorithm, image,
++ labels_onehot, labels, batch_rng, augmentation_on,
++ gradient_diagnostics)
++ elif learning_algorithm != "backprop":
+ state, metrics = train_step(
+ state, image, labels_onehot, labels, batch_rng, inf_rng,
+ augmentation_on, gradient_diagnostics)
+@@ -281,6 +315,58 @@ def train_epoch(state, train_ds, batch_size, rng, augmentation_on,
+
+ return state, epoch_metrics_np, runtime
+
++
++def direct_feedback_fields(states, output_field, direct_feedback):
++ """Apply per-hidden fixed output maps without a layerwise reverse chain."""
++ fields = [jnp.zeros_like(states[0])]
++ for feedback in direct_feedback:
++ fields.append(jnp.tensordot(
++ output_field, feedback, axes=((-1,), (0,))))
++ fields.append(output_field)
++ return fields
++
++
++@partial(jax.jit, static_argnames=("learning_algorithm",))
++def train_step_local(state, local_feedback, learning_algorithm, image,
++ labels_onehot, labels, batch_rng, augmentation_on,
++ gradient_diagnostics):
++ """One ordinary-FA or DFA step with detached local eligibilities."""
++ image = jax.lax.cond(
++ augmentation_on, vmap_augment_train, no_aug, image, batch_rng)
++ states, linear = state.apply_fn(
++ {"params": state.params}, image, method="ff_with_local_cache")
++ states = tree_map(jax.lax.stop_gradient, states)
++ linear = tree_map(jax.lax.stop_gradient, linear)
++
++ def output_loss(logits):
++ return state.apply_fn(
++ {"params": state.params}, logits, labels_onehot,
++ method="output_loss")
++
++ output_field = jax.lax.stop_gradient(jax.grad(output_loss)(states[-1]))
++ if learning_algorithm == "fa":
++ teaching_fields = state.apply_fn(
++ {"params": local_feedback}, states, linear, output_field,
++ method="fa_teaching_fields")
++ else:
++ teaching_fields = direct_feedback_fields(
++ states, output_field, local_feedback)
++ teaching_fields = tree_map(jax.lax.stop_gradient, teaching_fields)
++
++ def local_objective(params):
++ return state.apply_fn(
++ {"params": params}, states, teaching_fields,
++ method="local_correlation_objective")
++
++ grads = jax.grad(local_objective)(state.params)
++ metrics = compute_metrics(
++ image=image, labels_onehot=labels_onehot, labels=labels, state=state)
++ metrics = jax.lax.cond(
++ gradient_diagnostics, ref_grad_and_angle, no_ref_grad_and_angle,
++ state, grads, image, labels_onehot, metrics)
++ state = state.apply_gradients(grads=grads)
++ return state, metrics
++
+ @jax.jit
+ def train_step(state, image, labels_onehot, labels, batch_rng, inf_rng,
+ augmentation_on, gradient_diagnostics):
+diff --git a/tests/local_rules_smoke.py b/tests/local_rules_smoke.py
+new file mode 100644
+index 0000000..2366e61
+--- /dev/null
++++ b/tests/local_rules_smoke.py
+@@ -0,0 +1,138 @@
++#!/usr/bin/env python3
++"""Deterministic mechanics checks for matched ordinary FA and DFA."""
++import json
++import os
++import sys
++
++import jax
++import jax.numpy as jnp
++import optax
++from flax import linen as nn
++from flax.core.frozen_dict import unfreeze
++
++
++ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
++sys.path.insert(0, ROOT)
++
++from src.models import cnn_abstract
++from src.training_utils import create_local_feedback, direct_feedback_fields
++
++
++def loss_func(logits, one_hot):
++ return jnp.sum(optax.softmax_cross_entropy(logits, one_hot))
++
++
++def flat(tree):
++ return jax.flatten_util.ravel_pytree(tree)[0].astype(jnp.float64)
++
++
++def relative_error(left, right):
++ return float(
++ jnp.linalg.norm(flat(left) - flat(right))
++ / jnp.maximum(jnp.linalg.norm(flat(right)), 1e-30)
++ )
++
++
++def replace_layer(tree, layer_name, delta):
++ changed = unfreeze(tree)
++ changed[layer_name] = jax.tree_util.tree_map(
++ lambda value: value + delta, changed[layer_name])
++ return changed
++
++
++def main():
++ jax.config.update("jax_enable_x64", True)
++ model = cnn_abstract(
++ loss_func, nn.Conv, nn.Dense, nn.relu, 3, 0.1, 0.0,
++ jnp.float64, jnp.float64,
++ kernels=[(3, 3), (3, 3)], strides=[(1, 1), (1, 1)],
++ features=[4, 5], mp=[True, True], dense_features=[3],
++ inference_sequence="fwK", inference_passes_nudged=1)
++ x = jax.random.normal(jax.random.PRNGKey(1), (3, 8, 8, 2),
++ dtype=jnp.float64)
++ labels = jnp.asarray([0, 2, 1])
++ one_hot = jax.nn.one_hot(labels, 3, dtype=jnp.float64)
++ params = unfreeze(model.init(jax.random.PRNGKey(2), x)["params"])
++ states, linear = model.apply(
++ {"params": params}, x, method="ff_with_local_cache")
++
++ def batch_loss(candidate):
++ logits = model.apply({"params": candidate}, x)
++ return model.apply(
++ {"params": candidate}, logits, one_hot,
++ method="output_loss") / x.shape[0]
++
++ bp_grads = jax.grad(batch_loss)(params)
++ output_field = jax.grad(
++ lambda logits: model.apply(
++ {"params": params}, logits, one_hot, method="output_loss")
++ )(states[-1])
++ symmetric_fields = model.apply(
++ {"params": params}, states, linear, output_field,
++ method="fa_teaching_fields")
++ symmetric_grads = jax.grad(
++ lambda candidate: model.apply(
++ {"params": candidate}, states, symmetric_fields,
++ method="local_correlation_objective")
++ )(params)
++ symmetric_error = relative_error(symmetric_grads, bp_grads)
++ assert symmetric_error < 2e-12, symmetric_error
++
++ feedback = create_local_feedback(
++ jax.random.PRNGKey(1729), model, params, (8, 8, 2), "fa", 3)
++ feedback_cosine = float(
++ jnp.vdot(flat(feedback), flat(params))
++ / (jnp.linalg.norm(flat(feedback)) * jnp.linalg.norm(flat(params)))
++ )
++ assert abs(feedback_cosine) < 0.25, feedback_cosine
++ random_fields = model.apply(
++ {"params": feedback}, states, linear, output_field,
++ method="fa_teaching_fields")
++ changed_forward = replace_layer(params, "d00", 0.75)
++ random_fields_again = model.apply(
++ {"params": feedback}, states, linear, output_field,
++ method="fa_teaching_fields")
++ feedback_independence_error = relative_error(
++ random_fields[1:-1], random_fields_again[1:-1])
++ assert feedback_independence_error == 0.0
++
++ random_fields = jax.tree_util.tree_map(
++ jax.lax.stop_gradient, random_fields)
++ local_grads = jax.grad(
++ lambda candidate: model.apply(
++ {"params": candidate}, states, random_fields,
++ method="local_correlation_objective")
++ )(params)
++ changed_local_grads = jax.grad(
++ lambda candidate: model.apply(
++ {"params": candidate}, states, random_fields,
++ method="local_correlation_objective")
++ )(changed_forward)
++ local_boundary_error = relative_error(
++ local_grads["c00"], changed_local_grads["c00"])
++ assert local_boundary_error == 0.0
++
++ direct = create_local_feedback(
++ jax.random.PRNGKey(1730), model, params, (8, 8, 2), "dfa", 3)
++ assert len(direct) == len(states) - 2
++ dfa_fields = direct_feedback_fields(states, output_field, direct)
++ assert len(dfa_fields) == len(states)
++ assert all(
++ field.shape == state.shape
++ for field, state in zip(dfa_fields, states)
++ )
++ assert all(bool(jnp.all(jnp.isfinite(field))) for field in dfa_fields)
++
++ report = {
++ "symmetric_fa_bp_relative_error": symmetric_error,
++ "independent_feedback_forward_cosine": feedback_cosine,
++ "feedback_independence_error": feedback_independence_error,
++ "detached_local_boundary_error": local_boundary_error,
++ "dfa_hidden_maps": len(direct),
++ "status": "passed",
++ }
++ print(json.dumps(report, indent=2, sort_keys=True))
++
++
++if __name__ == "__main__":
++ main()
+diff --git a/train.py b/train.py
+index 139fd6c..c030158 100644
+--- a/train.py
++++ b/train.py
+@@ -8,7 +8,7 @@ from absl import logging # for logging
+ from matplotlib.colors import LogNorm
+
+ # Training utils
+-from src import create_train_state, train_epoch, eval_model, heatmap_grads_batches, heatmap_grads_epochs, plot_L_or_gamma
++from src import create_train_state, create_local_feedback, train_epoch, eval_model, heatmap_grads_batches, heatmap_grads_epochs, plot_L_or_gamma
+
+ # Import configurations
+ # import config # Use this for the old method
+@@ -46,6 +46,9 @@ for experiment_index, seed in enumerate(config.seeds):
+ steps_per_epoch = len(config.train_ds['image']) // config.batch_size
+
+ state = create_train_state(init_rng, config.model, config.image_dims, config.learning_rate, config.warmup_learning_rate, config.learning_rate_final, config.momentum, config.weight_decay, config.num_epochs, config.warmup_epochs, config.decay_epochs, steps_per_epoch)
++ local_feedback = create_local_feedback(
++ jax.random.PRNGKey(config.feedback_seed), config.model, state.params,
++ config.image_dims, config.learning_algorithm, config.num_classes)
+
+
+ del init_rng # Must not be used anymore.
+@@ -74,6 +77,7 @@ for experiment_index, seed in enumerate(config.seeds):
+ state, epoch_metrics, train_time = train_epoch(
+ state, config.train_ds, config.batch_size, input_rng,
+ augmentation_on, config.learning_algorithm, config.num_classes,
++ local_feedback=local_feedback,
+ gradient_diagnostics=(config.gradient_diagnostics == "full"))
+ loginfo_and_print('train: \tloss: %.4f, \taccuracy: %.4f, \truntime: %.4f' % (epoch_metrics["loss"], epoch_metrics["accuracy"], train_time))
+ hist['train_loss'][epoch-1], hist['train_accuracy'][epoch-1], hist['train_time'][epoch-1] = epoch_metrics["loss"], epoch_metrics["accuracy"], train_time
+--
+2.54.0
+
diff --git a/external/dualprop_patches/0003-fix-norm-preserve-direct-feedback-maps.patch b/external/dualprop_patches/0003-fix-norm-preserve-direct-feedback-maps.patch
new file mode 100644
index 0000000..d2af3e1
--- /dev/null
+++ b/external/dualprop_patches/0003-fix-norm-preserve-direct-feedback-maps.patch
@@ -0,0 +1,43 @@
+From ce106aa21a77ee92cc98ea543ea0fc4fe211078a Mon Sep 17 00:00:00 2001
+From: YurenHao0426 <Blackhao0426@gmail.com>
+Date: Mon, 27 Jul 2026 13:05:38 -0500
+Subject: [PATCH 03/19] fix: norm-preserve direct feedback maps
+
+---
+ src/training_utils.py | 20 +++++++++++++-------
+ 1 file changed, 13 insertions(+), 7 deletions(-)
+
+diff --git a/src/training_utils.py b/src/training_utils.py
+index b6e1e3c..2b6b18a 100644
+--- a/src/training_utils.py
++++ b/src/training_utils.py
+@@ -236,13 +236,19 @@ def create_local_feedback(rng, model, params, image_dims, learning_algorithm,
+ states, _ = model.apply(
+ {"params": params}, dummy, method="ff_with_local_cache")
+ keys = jax.random.split(rng, len(states) - 2)
+- scale = jnp.asarray(num_classes ** -0.5, dtype=dummy.dtype)
+- return tuple(
+- scale * jax.random.normal(
+- key, (num_classes,) + tuple(state.shape[1:]),
+- dtype=state.dtype)
+- for key, state in zip(keys, states[1:-1])
+- )
++ feedback = []
++ for key, state in zip(keys, states[1:-1]):
++ hidden_units = int(np.prod(state.shape[1:]))
++ # Preserve the norm of an output teaching vector in expectation. A
++ # 1/sqrt(num_classes) per-coordinate scale would make the *whole*
++ # hidden teaching field grow as sqrt(num_hidden_units), which is
++ # catastrophic for early convolutional feature maps.
++ scale = jnp.asarray(hidden_units ** -0.5, dtype=state.dtype)
++ feedback.append(
++ scale * jax.random.normal(
++ key, (num_classes,) + tuple(state.shape[1:]),
++ dtype=state.dtype))
++ return tuple(feedback)
+
+ def augment_train(image, batch_rng):
+ w, h, c = image.shape
+--
+2.54.0
+
diff --git a/external/dualprop_patches/0004-crossover-add-architecture-compatible-PEPITA.patch b/external/dualprop_patches/0004-crossover-add-architecture-compatible-PEPITA.patch
new file mode 100644
index 0000000..896350b
--- /dev/null
+++ b/external/dualprop_patches/0004-crossover-add-architecture-compatible-PEPITA.patch
@@ -0,0 +1,242 @@
+From 3bfbe71fcdc138764408a5e1646b66c464308b8b Mon Sep 17 00:00:00 2001
+From: YurenHao0426 <Blackhao0426@gmail.com>
+Date: Mon, 27 Jul 2026 13:09:05 -0500
+Subject: [PATCH 04/19] crossover: add architecture-compatible PEPITA
+
+---
+ config/cli_config.py | 8 ++++--
+ src/models.py | 37 ++++++++++++++++++++++++++++
+ src/training_utils.py | 50 ++++++++++++++++++++++++++++++++++++--
+ tests/local_rules_smoke.py | 32 ++++++++++++++++++++++++
+ train.py | 3 ++-
+ 5 files changed, 125 insertions(+), 5 deletions(-)
+
+diff --git a/config/cli_config.py b/config/cli_config.py
+index e1756e1..3088436 100644
+--- a/config/cli_config.py
++++ b/config/cli_config.py
+@@ -42,12 +42,16 @@ parser.add_argument('--experiment-name', default='test', help='A string denoting
+
+ parser.add_argument('--model', default='VGG16', choices=['VGG16', 'VGGlike', 'CNN', 'miniCNN', 'MLP'], help='')
+
+-parser.add_argument('--learning-algorithm', default='dualprop-lagr-ff', choices=['backprop', 'fa', 'dfa', 'dualprop-lagr-ff', 'dualprop-raovr-ff', 'dualprop-raovr-dampened-ff'])
++parser.add_argument('--learning-algorithm', default='dualprop-lagr-ff', choices=['backprop', 'fa', 'dfa', 'pepita', 'dualprop-lagr-ff', 'dualprop-raovr-ff', 'dualprop-raovr-dampened-ff'])
+
+ parser.add_argument(
+ '--feedback-seed', default=1729, type=int,
+ help='Independent fixed-feedback initialization seed for FA/DFA.')
+
++parser.add_argument(
++ '--pepita-projection-scale', default=0.05, type=float,
++ help='Multiplier on PEPITA He-uniform output-error-to-input projection.')
++
+ parser.add_argument(
+ '--gradient-diagnostics', default='full', choices=['none', 'full'],
+ help=('Compute the exact BP reference gradient and layerwise cosine on '
+@@ -128,7 +132,7 @@ elif config.model == "MLP":
+ dense_features = [1024, 1024, config.num_classes]
+
+ # Load model
+-modeltype = {"backprop":cnn_abstract, "fa": cnn_abstract, "dfa": cnn_abstract, "dualprop-lagr-ff": cnn_dualprop_Lagr_ff, "dualprop-raovr-ff": cnn_dualprop_RAOVR_ff, "dualprop-raovr-dampened-ff": cnn_dualprop_RAOVR_dampened_ff}
++modeltype = {"backprop":cnn_abstract, "fa": cnn_abstract, "dfa": cnn_abstract, "pepita": cnn_abstract, "dualprop-lagr-ff": cnn_dualprop_Lagr_ff, "dualprop-raovr-ff": cnn_dualprop_RAOVR_ff, "dualprop-raovr-dampened-ff": cnn_dualprop_RAOVR_dampened_ff}
+ activation={"relu": relu, "hs": hs, "sigmoid": sigmoid, "tanh": tanh}
+ config.model = modeltype[config.learning_algorithm](loss_func, Conv, Dense, activation[config.activation], config.num_classes, config.beta, config.alpha, config.dtype, config.param_dtype,
+ kernels=kernels, strides=strides, features=features, mp = mp,
+diff --git a/src/models.py b/src/models.py
+index edc432c..127268a 100644
+--- a/src/models.py
++++ b/src/models.py
+@@ -194,6 +194,43 @@ class cnn_abstract(nn.Module, ABC):
+ _, linear_pullback = jax.vjp(linear_map, s[i])
+ fields[i] = linear_pullback(linear_field)[0]
+ return fields
++
++ def pepita_correlation_objective(self, modulated_states, original_linear,
++ modulated_linear, one_hot):
++ """Architecture-compatible PEPITA/ERIN two-presentation update.
++
++ Hidden fields are the first-minus-second post-ReLU activity
++ differences from the PEPITA equation. They multiply the modulated
++ presynaptic state directly, without differentiating a nonlinearity.
++ Convolutional correlations follow the reference code's additional
++ average over spatial positions. The readout uses the second-pass
++ softmax error.
++ """
++ objective = 0.0
++ batch_size = modulated_states[0].shape[0]
++ for i, layer in enumerate(self.layers):
++ x = jax.lax.stop_gradient(modulated_states[i])
++ if i == self.num_layers - 1:
++ field = jax.lax.stop_gradient(
++ jax.nn.softmax(modulated_states[-1], axis=-1) - one_hot)
++ prediction = layer(x)
++ objective += jnp.sum(prediction * field) / batch_size
++ elif i < self.num_convlayers:
++ field = jax.lax.stop_gradient(
++ self.act(original_linear[i])
++ - self.act(modulated_linear[i]))
++ prediction = layer.call_without_pooling(x)
++ spatial_positions = prediction.shape[1] * prediction.shape[2]
++ objective += (
++ jnp.sum(prediction * field)
++ / (batch_size * spatial_positions))
++ else:
++ field = jax.lax.stop_gradient(
++ self.act(original_linear[i])
++ - self.act(modulated_linear[i]))
++ prediction = layer(x)
++ objective += jnp.sum(prediction * field) / batch_size
++ return objective
+
+ def init_states_to_zero(self, x0):
+ s = [x0]
+diff --git a/src/training_utils.py b/src/training_utils.py
+index 2b6b18a..f5fae62 100644
+--- a/src/training_utils.py
++++ b/src/training_utils.py
+@@ -219,17 +219,23 @@ def create_train_state(rng, model, image_dims, lr, wlr, lrf, momentum, weight_de
+
+
+ def create_local_feedback(rng, model, params, image_dims, learning_algorithm,
+- num_classes):
++ num_classes, pepita_projection_scale=0.05):
+ """Create fixed feedback without reading a forward parameter value.
+
+ FA uses an independently initialized model-shaped linear parameter tree.
+ DFA uses one independent output-to-hidden tensor per hidden population.
+ The forward tree is used only to infer activation shapes for DFA.
+ """
+- if learning_algorithm not in ("fa", "dfa"):
++ if learning_algorithm not in ("fa", "dfa", "pepita"):
+ return None
+ w, h, channels = image_dims
+ dummy = jnp.ones([1, w, h, channels])
++ if learning_algorithm == "pepita":
++ input_units = w * h * channels
++ limit = jnp.sqrt(6.0 / input_units) * pepita_projection_scale
++ return jax.random.uniform(
++ rng, (num_classes, w, h, channels),
++ minval=-limit, maxval=limit, dtype=dummy.dtype)
+ if learning_algorithm == "fa":
+ return unfreeze(model.init(rng, dummy)["params"])
+
+@@ -301,6 +307,10 @@ def train_epoch(state, train_ds, batch_size, rng, augmentation_on,
+ state, local_feedback, learning_algorithm, image,
+ labels_onehot, labels, batch_rng, augmentation_on,
+ gradient_diagnostics)
++ elif learning_algorithm == "pepita":
++ state, metrics = train_step_pepita(
++ state, local_feedback, image, labels_onehot, labels,
++ batch_rng, augmentation_on, gradient_diagnostics)
+ elif learning_algorithm != "backprop":
+ state, metrics = train_step(
+ state, image, labels_onehot, labels, batch_rng, inf_rng,
+@@ -332,6 +342,42 @@ def direct_feedback_fields(states, output_field, direct_feedback):
+ return fields
+
+
++@jax.jit
++def train_step_pepita(state, input_feedback, image, labels_onehot, labels,
++ batch_rng, augmentation_on, gradient_diagnostics):
++ """One PEPITA step: ordinary forward, error-modulated forward, local rule."""
++ image = jax.lax.cond(
++ augmentation_on, vmap_augment_train, no_aug, image, batch_rng)
++ states, linear = state.apply_fn(
++ {"params": state.params}, image, method="ff_with_local_cache")
++ output_error = jax.lax.stop_gradient(
++ jax.nn.softmax(states[-1], axis=-1) - labels_onehot)
++ input_error = jnp.tensordot(
++ output_error, input_feedback, axes=((-1,), (0,)))
++ modulated_image = jax.lax.stop_gradient(image + input_error)
++ modulated_states, modulated_linear = state.apply_fn(
++ {"params": state.params}, modulated_image,
++ method="ff_with_local_cache")
++ states = tree_map(jax.lax.stop_gradient, states)
++ linear = tree_map(jax.lax.stop_gradient, linear)
++ modulated_states = tree_map(jax.lax.stop_gradient, modulated_states)
++ modulated_linear = tree_map(jax.lax.stop_gradient, modulated_linear)
++
++ def local_objective(params):
++ return state.apply_fn(
++ {"params": params}, modulated_states, linear, modulated_linear,
++ labels_onehot, method="pepita_correlation_objective")
++
++ grads = jax.grad(local_objective)(state.params)
++ metrics = compute_metrics(
++ image=image, labels_onehot=labels_onehot, labels=labels, state=state)
++ metrics = jax.lax.cond(
++ gradient_diagnostics, ref_grad_and_angle, no_ref_grad_and_angle,
++ state, grads, image, labels_onehot, metrics)
++ state = state.apply_gradients(grads=grads)
++ return state, metrics
++
++
+ @partial(jax.jit, static_argnames=("learning_algorithm",))
+ def train_step_local(state, local_feedback, learning_algorithm, image,
+ labels_onehot, labels, batch_rng, augmentation_on,
+diff --git a/tests/local_rules_smoke.py b/tests/local_rules_smoke.py
+index 2366e61..e7da52e 100644
+--- a/tests/local_rules_smoke.py
++++ b/tests/local_rules_smoke.py
+@@ -123,12 +123,44 @@ def main():
+ )
+ assert all(bool(jnp.all(jnp.isfinite(field))) for field in dfa_fields)
+
++ pepita_projection = create_local_feedback(
++ jax.random.PRNGKey(1731), model, params, (8, 8, 2), "pepita", 3,
++ pepita_projection_scale=0.05)
++ pepita_error = jax.nn.softmax(states[-1], axis=-1) - one_hot
++ modulated_x = x + jnp.tensordot(
++ pepita_error, pepita_projection, axes=((-1,), (0,)))
++ modulated_states, modulated_linear = model.apply(
++ {"params": params}, modulated_x, method="ff_with_local_cache")
++ pepita_grads = jax.grad(
++ lambda candidate: model.apply(
++ {"params": candidate}, modulated_states, linear,
++ modulated_linear, one_hot,
++ method="pepita_correlation_objective")
++ )(params)
++ output_error = jax.nn.softmax(modulated_states[-1], axis=-1) - one_hot
++ output_input = modulated_states[-2].reshape((x.shape[0], -1))
++ manual_kernel = output_input.T @ output_error / x.shape[0]
++ manual_bias = output_error.mean(axis=0)
++ readout_leaves = jax.tree_util.tree_leaves(pepita_grads["d00"])
++ readout_kernel = next(value for value in readout_leaves if value.ndim == 2)
++ readout_bias = next(value for value in readout_leaves if value.ndim == 1)
++ pepita_readout_error = max(
++ float(jnp.max(jnp.abs(readout_kernel - manual_kernel))),
++ float(jnp.max(jnp.abs(readout_bias - manual_bias))),
++ )
++ assert pepita_readout_error < 2e-12, pepita_readout_error
++ projection_limit = (6.0 / (8 * 8 * 2)) ** 0.5 * 0.05
++ assert float(jnp.max(jnp.abs(pepita_projection))) <= projection_limit
++ assert bool(jnp.all(jnp.isfinite(flat(pepita_grads))))
++
+ report = {
+ "symmetric_fa_bp_relative_error": symmetric_error,
+ "independent_feedback_forward_cosine": feedback_cosine,
+ "feedback_independence_error": feedback_independence_error,
+ "detached_local_boundary_error": local_boundary_error,
+ "dfa_hidden_maps": len(direct),
++ "pepita_readout_equation_max_error": pepita_readout_error,
++ "pepita_projection_shape": list(pepita_projection.shape),
+ "status": "passed",
+ }
+ print(json.dumps(report, indent=2, sort_keys=True))
+diff --git a/train.py b/train.py
+index c030158..ba5f6ac 100644
+--- a/train.py
++++ b/train.py
+@@ -48,7 +48,8 @@ for experiment_index, seed in enumerate(config.seeds):
+ state = create_train_state(init_rng, config.model, config.image_dims, config.learning_rate, config.warmup_learning_rate, config.learning_rate_final, config.momentum, config.weight_decay, config.num_epochs, config.warmup_epochs, config.decay_epochs, steps_per_epoch)
+ local_feedback = create_local_feedback(
+ jax.random.PRNGKey(config.feedback_seed), config.model, state.params,
+- config.image_dims, config.learning_algorithm, config.num_classes)
++ config.image_dims, config.learning_algorithm, config.num_classes,
++ pepita_projection_scale=config.pepita_projection_scale)
+
+
+ del init_rng # Must not be used anymore.
+--
+2.54.0
+
diff --git a/external/dualprop_patches/0005-crossover-add-greedy-Forward-Forward-runner.patch b/external/dualprop_patches/0005-crossover-add-greedy-Forward-Forward-runner.patch
new file mode 100644
index 0000000..1b5f46f
--- /dev/null
+++ b/external/dualprop_patches/0005-crossover-add-greedy-Forward-Forward-runner.patch
@@ -0,0 +1,422 @@
+From fe79d6281891ac52977416bb9d65faac3c32e8e4 Mon Sep 17 00:00:00 2001
+From: YurenHao0426 <Blackhao0426@gmail.com>
+Date: Mon, 27 Jul 2026 13:12:55 -0500
+Subject: [PATCH 05/19] crossover: add greedy Forward-Forward runner
+
+---
+ config/cli_config.py | 12 ++++-
+ src/__init__.py | 2 +-
+ src/models.py | 43 +++++++++++++++
+ src/training_utils.py | 108 +++++++++++++++++++++++++++++++++++++
+ tests/local_rules_smoke.py | 40 +++++++++++++-
+ train.py | 4 ++
+ train_ff.py | 92 +++++++++++++++++++++++++++++++
+ 7 files changed, 297 insertions(+), 4 deletions(-)
+ create mode 100644 train_ff.py
+
+diff --git a/config/cli_config.py b/config/cli_config.py
+index 3088436..159e9bc 100644
+--- a/config/cli_config.py
++++ b/config/cli_config.py
+@@ -42,7 +42,7 @@ parser.add_argument('--experiment-name', default='test', help='A string denoting
+
+ parser.add_argument('--model', default='VGG16', choices=['VGG16', 'VGGlike', 'CNN', 'miniCNN', 'MLP'], help='')
+
+-parser.add_argument('--learning-algorithm', default='dualprop-lagr-ff', choices=['backprop', 'fa', 'dfa', 'pepita', 'dualprop-lagr-ff', 'dualprop-raovr-ff', 'dualprop-raovr-dampened-ff'])
++parser.add_argument('--learning-algorithm', default='dualprop-lagr-ff', choices=['backprop', 'fa', 'dfa', 'pepita', 'ff', 'dualprop-lagr-ff', 'dualprop-raovr-ff', 'dualprop-raovr-dampened-ff'])
+
+ parser.add_argument(
+ '--feedback-seed', default=1729, type=int,
+@@ -52,6 +52,14 @@ parser.add_argument(
+ '--pepita-projection-scale', default=0.05, type=float,
+ help='Multiplier on PEPITA He-uniform output-error-to-input projection.')
+
++parser.add_argument(
++ '--ff-threshold', default=2.0, type=float,
++ help='Forward-Forward positive/negative goodness threshold.')
++
++parser.add_argument(
++ '--ff-score-from-layer', default=1, type=int,
++ help='First zero-indexed FF layer included in candidate-label goodness.')
++
+ parser.add_argument(
+ '--gradient-diagnostics', default='full', choices=['none', 'full'],
+ help=('Compute the exact BP reference gradient and layerwise cosine on '
+@@ -132,7 +140,7 @@ elif config.model == "MLP":
+ dense_features = [1024, 1024, config.num_classes]
+
+ # Load model
+-modeltype = {"backprop":cnn_abstract, "fa": cnn_abstract, "dfa": cnn_abstract, "pepita": cnn_abstract, "dualprop-lagr-ff": cnn_dualprop_Lagr_ff, "dualprop-raovr-ff": cnn_dualprop_RAOVR_ff, "dualprop-raovr-dampened-ff": cnn_dualprop_RAOVR_dampened_ff}
++modeltype = {"backprop":cnn_abstract, "fa": cnn_abstract, "dfa": cnn_abstract, "pepita": cnn_abstract, "ff": cnn_abstract, "dualprop-lagr-ff": cnn_dualprop_Lagr_ff, "dualprop-raovr-ff": cnn_dualprop_RAOVR_ff, "dualprop-raovr-dampened-ff": cnn_dualprop_RAOVR_dampened_ff}
+ activation={"relu": relu, "hs": hs, "sigmoid": sigmoid, "tanh": tanh}
+ config.model = modeltype[config.learning_algorithm](loss_func, Conv, Dense, activation[config.activation], config.num_classes, config.beta, config.alpha, config.dtype, config.param_dtype,
+ kernels=kernels, strides=strides, features=features, mp = mp,
+diff --git a/src/__init__.py b/src/__init__.py
+index 7a08c87..752dcc2 100644
+--- a/src/__init__.py
++++ b/src/__init__.py
+@@ -1,2 +1,2 @@
+ from .models import cnn_dualprop_Lagr_ff, cnn_dualprop_RAOVR_ff, cnn_dualprop_RAOVR_dampened_ff, cnn_abstract
+-from .training_utils import create_train_state, create_local_feedback, train_epoch, eval_model, get_mnist, get_svhn, get_fashionmnist, get_cifar10, get_cifar100, get_imagenet_32x32, heatmap_grads_batches, heatmap_grads_epochs, plot_L_or_gamma
++from .training_utils import create_train_state, create_ff_train_state, create_local_feedback, train_epoch, train_ff_epoch, eval_model, eval_ff_model, get_mnist, get_svhn, get_fashionmnist, get_cifar10, get_cifar100, get_imagenet_32x32, heatmap_grads_batches, heatmap_grads_epochs, plot_L_or_gamma
+diff --git a/src/models.py b/src/models.py
+index 127268a..0f5af4f 100644
+--- a/src/models.py
++++ b/src/models.py
+@@ -231,6 +231,49 @@ class cnn_abstract(nn.Module, ABC):
+ prediction = layer(x)
+ objective += jnp.sum(prediction * field) / batch_size
+ return objective
++
++ @staticmethod
++ def _ff_normalize(x):
++ axes = tuple(range(1, x.ndim))
++ norm = jnp.sqrt(jnp.sum(jnp.square(x), axis=axes, keepdims=True))
++ return x / (norm + 1e-8)
++
++ def _ff_layer(self, x, layer_index):
++ x = self._ff_normalize(x)
++ return self.act(self.layers[layer_index](x))
++
++ def ff_prefix(self, x, stop_layer):
++ """Detached input producer for greedy Forward-Forward training."""
++ for i in range(stop_layer):
++ x = self._ff_layer(x, i)
++ return x
++
++ def ff_layer_objective(self, positive_input, negative_input, layer_index,
++ threshold):
++ """Reference Forward-Forward local softplus goodness objective."""
++ positive = self._ff_layer(
++ jax.lax.stop_gradient(positive_input), layer_index)
++ negative = self._ff_layer(
++ jax.lax.stop_gradient(negative_input), layer_index)
++ axes = tuple(range(1, positive.ndim))
++ positive_goodness = jnp.mean(jnp.square(positive), axis=axes)
++ negative_goodness = jnp.mean(jnp.square(negative), axis=axes)
++ loss = jnp.mean(
++ jax.nn.softplus(-positive_goodness + threshold)
++ + jax.nn.softplus(negative_goodness - threshold))
++ pair_accuracy = jnp.mean(
++ positive_goodness > negative_goodness, dtype=jnp.float32)
++ return loss, (positive_goodness, negative_goodness, pair_accuracy)
++
++ def ff_goodness(self, x, score_from_layer=1):
++ """Candidate-label goodness used by supervised FF inference."""
++ score = jnp.zeros((x.shape[0],), dtype=x.dtype)
++ for i in range(self.num_layers):
++ x = self._ff_layer(x, i)
++ if i >= score_from_layer:
++ axes = tuple(range(1, x.ndim))
++ score += jnp.mean(jnp.square(x), axis=axes)
++ return score
+
+ def init_states_to_zero(self, x0):
+ s = [x0]
+diff --git a/src/training_utils.py b/src/training_utils.py
+index f5fae62..734f46d 100644
+--- a/src/training_utils.py
++++ b/src/training_utils.py
+@@ -218,6 +218,15 @@ def create_train_state(rng, model, image_dims, lr, wlr, lrf, momentum, weight_de
+ return train_state.TrainState.create(apply_fn=model.apply, params=unfreeze(params), tx=tx)
+
+
++def create_ff_train_state(rng, model, image_dims, learning_rate):
++ """Reference-style Adam state for greedy Forward-Forward layers."""
++ w, h, channels = image_dims
++ dummy = jnp.ones([1, w, h, channels])
++ params = unfreeze(model.init(rng, dummy)["params"])
++ return train_state.TrainState.create(
++ apply_fn=model.apply, params=params, tx=optax.adam(learning_rate))
++
++
+ def create_local_feedback(rng, model, params, image_dims, learning_algorithm,
+ num_classes, pepita_projection_scale=0.05):
+ """Create fixed feedback without reading a forward parameter value.
+@@ -342,6 +351,105 @@ def direct_feedback_fields(states, output_field, direct_feedback):
+ return fields
+
+
++def ff_overlay(image, labels, num_classes):
++ """Overlay a candidate label on the first input coordinates."""
++ flat = image.reshape((image.shape[0], -1))
++ flat = flat.at[:, :num_classes].set(0.0)
++ flat = flat.at[jnp.arange(image.shape[0]), labels].set(jnp.max(image))
++ return flat.reshape(image.shape)
++
++
++@partial(jax.jit, static_argnames=("layer_index", "num_classes"))
++def train_step_ff(state, image, labels, batch_rng, augmentation_on,
++ layer_index, num_classes, threshold):
++ """One greedy Forward-Forward update of exactly one layer."""
++ augmentation_rng, negative_rng = jax.random.split(batch_rng)
++ per_example_rng = jax.random.split(augmentation_rng, image.shape[0])
++ image = jax.lax.cond(
++ augmentation_on, vmap_augment_train, no_aug, image, per_example_rng)
++ offsets = jax.random.randint(
++ negative_rng, labels.shape, 1, num_classes)
++ negative_labels = (labels + offsets) % num_classes
++ positive = ff_overlay(image, labels, num_classes)
++ negative = ff_overlay(image, negative_labels, num_classes)
++ positive_input = state.apply_fn(
++ {"params": state.params}, positive, layer_index, method="ff_prefix")
++ negative_input = state.apply_fn(
++ {"params": state.params}, negative, layer_index, method="ff_prefix")
++ positive_input = jax.lax.stop_gradient(positive_input)
++ negative_input = jax.lax.stop_gradient(negative_input)
++
++ def objective(params):
++ return state.apply_fn(
++ {"params": params}, positive_input, negative_input, layer_index,
++ threshold, method="ff_layer_objective")
++
++ (loss, auxiliary), grads = jax.value_and_grad(
++ objective, has_aux=True)(state.params)
++ state = state.apply_gradients(grads=grads)
++ return state, {
++ "loss": loss,
++ "positive_goodness": jnp.mean(auxiliary[0]),
++ "negative_goodness": jnp.mean(auxiliary[1]),
++ "pair_accuracy": auxiliary[2],
++ }
++
++
++def train_ff_epoch(state, train_ds, batch_size, rng, augmentation_on,
++ layer_index, num_classes, threshold):
++ """Train one FF layer for one full data epoch."""
++ t0 = time.time()
++ train_ds_size = len(train_ds["image"])
++ steps_per_epoch = train_ds_size // batch_size
++ perms = jax.random.permutation(rng, train_ds_size)
++ perms = perms[:steps_per_epoch * batch_size]
++ perms = perms.reshape((steps_per_epoch, batch_size))
++ metrics = []
++ for permutation in perms:
++ image = train_ds["image"][permutation]
++ labels = train_ds["label"][permutation]
++ rng, batch_rng = jax.random.split(rng)
++ state, batch_metrics = train_step_ff(
++ state, image, labels, batch_rng, augmentation_on, layer_index,
++ num_classes, threshold)
++ metrics.append(batch_metrics)
++ host_metrics = jax.device_get(metrics)
++ summary = {
++ key: float(np.mean([record[key] for record in host_metrics]))
++ for key in host_metrics[0]
++ }
++ return state, summary, time.time() - t0
++
++
++@partial(jax.jit, static_argnames=("num_classes", "score_from_layer"))
++def predict_ff(state, image, num_classes, score_from_layer):
++ scores = []
++ for candidate in range(num_classes):
++ labels = jnp.full((image.shape[0],), candidate, dtype=jnp.int32)
++ overlaid = ff_overlay(image, labels, num_classes)
++ scores.append(state.apply_fn(
++ {"params": state.params}, overlaid, score_from_layer,
++ method="ff_goodness"))
++ return jnp.stack(scores, axis=-1)
++
++
++def eval_ff_model(state, dataset, batch_size, num_classes, score_from_layer):
++ """Evaluate all candidate-label overlays; no classifier head is assumed."""
++ t0 = time.time()
++ size = len(dataset["image"])
++ steps = size // batch_size
++ indices = jnp.arange(steps * batch_size).reshape((steps, batch_size))
++ correct = 0
++ total = 0
++ for index in indices:
++ image = dataset["image"][index]
++ labels = dataset["label"][index]
++ scores = predict_ff(state, image, num_classes, score_from_layer)
++ correct += int(jax.device_get(jnp.sum(jnp.argmax(scores, -1) == labels)))
++ total += labels.shape[0]
++ return 100.0 * correct / total, time.time() - t0
++
++
+ @jax.jit
+ def train_step_pepita(state, input_feedback, image, labels_onehot, labels,
+ batch_rng, augmentation_on, gradient_diagnostics):
+diff --git a/tests/local_rules_smoke.py b/tests/local_rules_smoke.py
+index e7da52e..8ce1e02 100644
+--- a/tests/local_rules_smoke.py
++++ b/tests/local_rules_smoke.py
+@@ -15,7 +15,11 @@ ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+ sys.path.insert(0, ROOT)
+
+ from src.models import cnn_abstract
+-from src.training_utils import create_local_feedback, direct_feedback_fields
++from src.training_utils import (
++ create_local_feedback,
++ direct_feedback_fields,
++ ff_overlay,
++)
+
+
+ def loss_func(logits, one_hot):
+@@ -153,6 +157,38 @@ def main():
+ assert float(jnp.max(jnp.abs(pepita_projection))) <= projection_limit
+ assert bool(jnp.all(jnp.isfinite(flat(pepita_grads))))
+
++ positive = ff_overlay(x, labels, 3)
++ negative = ff_overlay(x, (labels + 1) % 3, 3)
++ assert bool(jnp.all(positive.reshape((x.shape[0], -1))[
++ jnp.arange(x.shape[0]), labels] == jnp.max(x)))
++ positive_prefix = model.apply(
++ {"params": params}, positive, 1, method="ff_prefix")
++ negative_prefix = model.apply(
++ {"params": params}, negative, 1, method="ff_prefix")
++ (ff_loss, ff_aux), ff_grads = jax.value_and_grad(
++ lambda candidate: model.apply(
++ {"params": candidate}, positive_prefix, negative_prefix, 1, 2.0,
++ method="ff_layer_objective"),
++ has_aux=True,
++ )(params)
++ ff_manual_loss = jnp.mean(
++ jax.nn.softplus(-ff_aux[0] + 2.0)
++ + jax.nn.softplus(ff_aux[1] - 2.0))
++ ff_objective_error = float(jnp.abs(ff_loss - ff_manual_loss))
++ assert ff_objective_error < 2e-12
++ assert float(jnp.linalg.norm(flat(ff_grads["c01"]))) > 0.0
++ assert float(jnp.linalg.norm(flat(ff_grads["c00"]))) == 0.0
++ assert float(jnp.linalg.norm(flat(ff_grads["d00"]))) == 0.0
++ ff_scores = jnp.stack([
++ model.apply(
++ {"params": params},
++ ff_overlay(x, jnp.full(labels.shape, candidate), 3), 1,
++ method="ff_goodness")
++ for candidate in range(3)
++ ], axis=-1)
++ assert ff_scores.shape == (x.shape[0], 3)
++ assert bool(jnp.all(jnp.isfinite(ff_scores)))
++
+ report = {
+ "symmetric_fa_bp_relative_error": symmetric_error,
+ "independent_feedback_forward_cosine": feedback_cosine,
+@@ -161,6 +197,8 @@ def main():
+ "dfa_hidden_maps": len(direct),
+ "pepita_readout_equation_max_error": pepita_readout_error,
+ "pepita_projection_shape": list(pepita_projection.shape),
++ "ff_local_objective_error": ff_objective_error,
++ "ff_score_shape": list(ff_scores.shape),
+ "status": "passed",
+ }
+ print(json.dumps(report, indent=2, sort_keys=True))
+diff --git a/train.py b/train.py
+index ba5f6ac..d3f0de6 100644
+--- a/train.py
++++ b/train.py
+@@ -14,6 +14,10 @@ from src import create_train_state, create_local_feedback, train_epoch, eval_mod
+ # import config # Use this for the old method
+ from config.cli_config import config
+
++if config.learning_algorithm == "ff":
++ raise ValueError(
++ "Forward-Forward uses greedy layerwise training; run train_ff.py")
++
+ experiment_dir = "./runs/" + config.experiment_name + "/"
+ if experiment_dir == "./runs/debug-test/" and os.path.isdir(experiment_dir):
+ shutil.rmtree(experiment_dir)
+diff --git a/train_ff.py b/train_ff.py
+new file mode 100644
+index 0000000..3839dd3
+--- /dev/null
++++ b/train_ff.py
+@@ -0,0 +1,92 @@
++"""Greedy supervised Forward-Forward on the author plain-CNN topologies."""
++import datetime
++import os
++import time
++
++import jax
++import numpy as np
++
++from config.cli_config import config
++from src import create_ff_train_state, eval_ff_model, train_ff_epoch
++
++
++if config.learning_algorithm != "ff":
++ raise ValueError("train_ff.py requires --learning-algorithm ff")
++if config.ff_score_from_layer < 0:
++ raise ValueError("--ff-score-from-layer must be nonnegative")
++
++experiment_dir = os.path.join("runs", config.experiment_name)
++if os.path.isdir(experiment_dir):
++ raise FileExistsError(
++ "experiment directory exists; refusing to overwrite "
++ + experiment_dir)
++os.makedirs(experiment_dir)
++
++for experiment_index, seed in enumerate(config.seeds):
++ timestamp = datetime.datetime.fromtimestamp(time.time())
++ outpath = os.path.join(
++ experiment_dir, timestamp.strftime("%Y_%m_%d_%H_%M_%S"))
++ os.makedirs(outpath)
++ print(
++ f"Starting FF experiment {experiment_index + 1}/"
++ f"{len(config.seeds)} seed={seed}",
++ flush=True)
++ rng = jax.random.PRNGKey(seed)
++ rng, init_rng = jax.random.split(rng)
++ state = create_ff_train_state(
++ init_rng, config.model, config.image_dims, config.learning_rate)
++ num_layers = len(state.params)
++ if config.ff_score_from_layer >= num_layers:
++ raise ValueError("--ff-score-from-layer excludes every layer")
++ augmentation_on = config.dataset not in ("mnist", "fashionmnist")
++ history = {
++ "method": "ff",
++ "epochs_per_layer": config.num_epochs,
++ "num_layers": num_layers,
++ "threshold": config.ff_threshold,
++ "score_from_layer": config.ff_score_from_layer,
++ "learning_rate": config.learning_rate,
++ "layers": [],
++ }
++ started = time.time()
++ for layer_index in range(num_layers):
++ layer_record = {"layer": layer_index, "epochs": []}
++ for epoch in range(config.num_epochs):
++ rng, epoch_rng = jax.random.split(rng)
++ state, metrics, runtime = train_ff_epoch(
++ state, config.train_ds, config.batch_size, epoch_rng,
++ augmentation_on, layer_index, config.num_classes,
++ config.ff_threshold)
++ record = {
++ "epoch": epoch + 1,
++ "runtime": runtime,
++ **metrics,
++ }
++ layer_record["epochs"].append(record)
++ print(
++ f"layer={layer_index + 1}/{num_layers} "
++ f"epoch={epoch + 1}/{config.num_epochs} "
++ f"loss={metrics['loss']:.5f} "
++ f"pair_acc={100 * metrics['pair_accuracy']:.2f}% "
++ f"runtime={runtime:.3f}s",
++ flush=True)
++ history["layers"].append(layer_record)
++
++ validation_accuracy, validation_time = eval_ff_model(
++ state, config.val_ds, config.batch_size, config.num_classes,
++ config.ff_score_from_layer)
++ test_accuracy, test_time = eval_ff_model(
++ state, config.test_ds, config.batch_size, config.num_classes,
++ config.ff_score_from_layer)
++ history["final"] = {
++ "validation_accuracy": validation_accuracy,
++ "validation_time": validation_time,
++ "test_accuracy": test_accuracy,
++ "test_time": test_time,
++ "train_and_eval_wall": time.time() - started,
++ }
++ print(
++ f"final val_accuracy={validation_accuracy:.3f}% "
++ f"test_accuracy={test_accuracy:.3f}%",
++ flush=True)
++ np.save(os.path.join(outpath, "hist.npy"), history)
+--
+2.54.0
+
diff --git a/external/dualprop_patches/0006-crossover-add-two-phase-equilibrium-propagation.patch b/external/dualprop_patches/0006-crossover-add-two-phase-equilibrium-propagation.patch
new file mode 100644
index 0000000..1d30ef2
--- /dev/null
+++ b/external/dualprop_patches/0006-crossover-add-two-phase-equilibrium-propagation.patch
@@ -0,0 +1,377 @@
+From b925ba7f07389054cba7b90d2b0db3d7e1699b94 Mon Sep 17 00:00:00 2001
+From: YurenHao0426 <Blackhao0426@gmail.com>
+Date: Mon, 27 Jul 2026 13:16:22 -0500
+Subject: [PATCH 06/19] crossover: add two-phase equilibrium propagation
+
+---
+ config/cli_config.py | 20 +++++++-
+ src/__init__.py | 2 +-
+ src/models.py | 57 +++++++++++++++++++++++
+ src/training_utils.py | 93 ++++++++++++++++++++++++++++++++++++++
+ tests/local_rules_smoke.py | 34 ++++++++++++++
+ train.py | 30 ++++++++----
+ 6 files changed, 224 insertions(+), 12 deletions(-)
+
+diff --git a/config/cli_config.py b/config/cli_config.py
+index 159e9bc..795e6a3 100644
+--- a/config/cli_config.py
++++ b/config/cli_config.py
+@@ -42,7 +42,7 @@ parser.add_argument('--experiment-name', default='test', help='A string denoting
+
+ parser.add_argument('--model', default='VGG16', choices=['VGG16', 'VGGlike', 'CNN', 'miniCNN', 'MLP'], help='')
+
+-parser.add_argument('--learning-algorithm', default='dualprop-lagr-ff', choices=['backprop', 'fa', 'dfa', 'pepita', 'ff', 'dualprop-lagr-ff', 'dualprop-raovr-ff', 'dualprop-raovr-dampened-ff'])
++parser.add_argument('--learning-algorithm', default='dualprop-lagr-ff', choices=['backprop', 'fa', 'dfa', 'pepita', 'ff', 'ep', 'dualprop-lagr-ff', 'dualprop-raovr-ff', 'dualprop-raovr-dampened-ff'])
+
+ parser.add_argument(
+ '--feedback-seed', default=1729, type=int,
+@@ -60,6 +60,22 @@ parser.add_argument(
+ '--ff-score-from-layer', default=1, type=int,
+ help='First zero-indexed FF layer included in candidate-label goodness.')
+
++parser.add_argument(
++ '--ep-beta', default=0.5, type=float,
++ help='Magnitude of the randomly signed EP output nudge.')
++
++parser.add_argument(
++ '--ep-dt', default=0.5, type=float,
++ help='Euler step size for EP state relaxation.')
++
++parser.add_argument(
++ '--ep-free-steps', default=20, type=int,
++ help='Number of free-phase EP relaxation steps.')
++
++parser.add_argument(
++ '--ep-nudge-steps', default=4, type=int,
++ help='Number of nudged-phase EP relaxation steps.')
++
+ parser.add_argument(
+ '--gradient-diagnostics', default='full', choices=['none', 'full'],
+ help=('Compute the exact BP reference gradient and layerwise cosine on '
+@@ -140,7 +156,7 @@ elif config.model == "MLP":
+ dense_features = [1024, 1024, config.num_classes]
+
+ # Load model
+-modeltype = {"backprop":cnn_abstract, "fa": cnn_abstract, "dfa": cnn_abstract, "pepita": cnn_abstract, "ff": cnn_abstract, "dualprop-lagr-ff": cnn_dualprop_Lagr_ff, "dualprop-raovr-ff": cnn_dualprop_RAOVR_ff, "dualprop-raovr-dampened-ff": cnn_dualprop_RAOVR_dampened_ff}
++modeltype = {"backprop":cnn_abstract, "fa": cnn_abstract, "dfa": cnn_abstract, "pepita": cnn_abstract, "ff": cnn_abstract, "ep": cnn_abstract, "dualprop-lagr-ff": cnn_dualprop_Lagr_ff, "dualprop-raovr-ff": cnn_dualprop_RAOVR_ff, "dualprop-raovr-dampened-ff": cnn_dualprop_RAOVR_dampened_ff}
+ activation={"relu": relu, "hs": hs, "sigmoid": sigmoid, "tanh": tanh}
+ config.model = modeltype[config.learning_algorithm](loss_func, Conv, Dense, activation[config.activation], config.num_classes, config.beta, config.alpha, config.dtype, config.param_dtype,
+ kernels=kernels, strides=strides, features=features, mp = mp,
+diff --git a/src/__init__.py b/src/__init__.py
+index 752dcc2..58c9adc 100644
+--- a/src/__init__.py
++++ b/src/__init__.py
+@@ -1,2 +1,2 @@
+ from .models import cnn_dualprop_Lagr_ff, cnn_dualprop_RAOVR_ff, cnn_dualprop_RAOVR_dampened_ff, cnn_abstract
+-from .training_utils import create_train_state, create_ff_train_state, create_local_feedback, train_epoch, train_ff_epoch, eval_model, eval_ff_model, get_mnist, get_svhn, get_fashionmnist, get_cifar10, get_cifar100, get_imagenet_32x32, heatmap_grads_batches, heatmap_grads_epochs, plot_L_or_gamma
++from .training_utils import create_train_state, create_ff_train_state, create_local_feedback, train_epoch, train_ff_epoch, eval_model, eval_ep_model, eval_ff_model, get_mnist, get_svhn, get_fashionmnist, get_cifar10, get_cifar100, get_imagenet_32x32, heatmap_grads_batches, heatmap_grads_epochs, plot_L_or_gamma
+diff --git a/src/models.py b/src/models.py
+index 0f5af4f..68e9804 100644
+--- a/src/models.py
++++ b/src/models.py
+@@ -274,6 +274,63 @@ class cnn_abstract(nn.Module, ABC):
+ axes = tuple(range(1, x.ndim))
+ score += jnp.mean(jnp.square(x), axis=axes)
+ return score
++
++ @staticmethod
++ def _ep_rho(state):
++ return jnp.clip(state, 0.0, 1.0)
++
++ @staticmethod
++ def _ep_rhop(state):
++ return jnp.asarray(
++ (state >= 0.0) & (state <= 1.0), dtype=state.dtype)
++
++ def ep_relax(self, x, one_hot, beta, steps, dt, initial_states=None):
++ """Canonical hard-sigmoid EP dynamics on the layered author graph.
++
++ The update follows the original two-phase implementation: synchronous
++ leaky state dynamics, symmetric top-down interactions through the
++ forward weights, and a squared-error output nudge. Max-pool pullbacks
++ use the local switches induced by the current lower state.
++ """
++ if initial_states is None:
++ states = self.init_states_to_zero(x)
++ else:
++ states = initial_states
++ for _ in range(steps):
++ updated = [x]
++ for i in range(1, self.num_layers):
++ below = self.layers[i - 1](self._ep_rho(states[i - 1]))
++ above = grad(self.get_phi, argnums=(1))(
++ self._ep_rho(states[i + 1]),
++ self._ep_rho(states[i]),
++ self.layers[i])
++ drive = (
++ -self._ep_rho(states[i]) + below + above)
++ next_state = states[i] + dt * self._ep_rhop(states[i]) * drive
++ updated.append(self._ep_rho(next_state))
++
++ below = self.layers[-1](self._ep_rho(states[-2]))
++ drive = -self._ep_rho(states[-1]) + below
++ drive += 2.0 * beta * (one_hot - states[-1])
++ next_output = (
++ states[-1] + dt * self._ep_rhop(states[-1]) * drive)
++ updated.append(self._ep_rho(next_output))
++ states = updated
++ return states
++
++ def ep_contrastive_objective(self, free_states, nudged_states, beta):
++ """Local EP correlation difference for optimizer-style descent."""
++ free_phi = 0.0
++ nudged_phi = 0.0
++ for i, layer in enumerate(self.layers):
++ free_phi += self.get_phi(
++ self._ep_rho(free_states[i + 1]),
++ self._ep_rho(free_states[i]), layer)
++ nudged_phi += self.get_phi(
++ self._ep_rho(nudged_states[i + 1]),
++ self._ep_rho(nudged_states[i]), layer)
++ return (free_phi - nudged_phi) / (
++ beta * free_states[0].shape[0])
+
+ def init_states_to_zero(self, x0):
+ s = [x0]
+diff --git a/src/training_utils.py b/src/training_utils.py
+index 734f46d..f767b80 100644
+--- a/src/training_utils.py
++++ b/src/training_utils.py
+@@ -291,6 +291,7 @@ def to_float32(ptree):
+
+ def train_epoch(state, train_ds, batch_size, rng, augmentation_on,
+ learning_algorithm, num_classes, local_feedback=None,
++ ep_beta=0.5, ep_free_steps=20, ep_nudge_steps=4, ep_dt=0.5,
+ gradient_diagnostics=True):
+ """Train for a single epoch."""
+ t0 = time.time()
+@@ -320,6 +321,11 @@ def train_epoch(state, train_ds, batch_size, rng, augmentation_on,
+ state, metrics = train_step_pepita(
+ state, local_feedback, image, labels_onehot, labels,
+ batch_rng, augmentation_on, gradient_diagnostics)
++ elif learning_algorithm == "ep":
++ state, metrics = train_step_ep(
++ state, image, labels_onehot, labels, batch_rng, inf_rng,
++ augmentation_on, ep_beta, ep_free_steps, ep_nudge_steps,
++ ep_dt, gradient_diagnostics)
+ elif learning_algorithm != "backprop":
+ state, metrics = train_step(
+ state, image, labels_onehot, labels, batch_rng, inf_rng,
+@@ -351,6 +357,47 @@ def direct_feedback_fields(states, output_field, direct_feedback):
+ return fields
+
+
++@partial(
++ jax.jit,
++ static_argnames=("free_steps", "nudge_steps"),
++)
++def train_step_ep(state, image, labels_onehot, labels, batch_rng, inf_rng,
++ augmentation_on, beta, free_steps, nudge_steps, dt,
++ gradient_diagnostics):
++ """One canonical two-phase equilibrium-propagation update."""
++ image = jax.lax.cond(
++ augmentation_on, vmap_augment_train, no_aug, image, batch_rng)
++ free_states = state.apply_fn(
++ {"params": state.params}, image, labels_onehot, 0.0, free_steps, dt,
++ method="ep_relax")
++ beta_sign = jnp.where(
++ jax.random.bernoulli(inf_rng), 1.0, -1.0)
++ signed_beta = beta * beta_sign
++ nudged_states = state.apply_fn(
++ {"params": state.params}, image, labels_onehot, signed_beta,
++ nudge_steps, dt, free_states, method="ep_relax")
++ free_states = tree_map(jax.lax.stop_gradient, free_states)
++ nudged_states = tree_map(jax.lax.stop_gradient, nudged_states)
++
++ def contrastive_objective(params):
++ return state.apply_fn(
++ {"params": params}, free_states, nudged_states, signed_beta,
++ method="ep_contrastive_objective")
++
++ grads = jax.grad(contrastive_objective)(state.params)
++ logits = free_states[-1]
++ metrics = {
++ "loss": jnp.mean(jnp.square(logits - labels_onehot)),
++ "accuracy": 100.0 * jnp.mean(
++ jnp.argmax(logits, -1) == labels, dtype=jnp.float32),
++ }
++ metrics = jax.lax.cond(
++ gradient_diagnostics, ref_grad_and_angle, no_ref_grad_and_angle,
++ state, grads, image, labels_onehot, metrics)
++ state = state.apply_gradients(grads=grads)
++ return state, metrics
++
++
+ def ff_overlay(image, labels, num_classes):
+ """Overlay a candidate label on the first input coordinates."""
+ flat = image.reshape((image.shape[0], -1))
+@@ -637,6 +684,52 @@ def eval_model(state, params, test_ds, batch_size, num_classes, eval_rng,
+
+ return summary['loss'], summary['accuracy'], summary['top5accuracy'], runtime, L10, L20, gamma10, gamma20
+
++
++@partial(jax.jit, static_argnames=("free_steps",))
++def eval_step_ep(state, params, image, labels_onehot, labels, free_steps, dt):
++ states = state.apply_fn(
++ {"params": params}, image, labels_onehot, 0.0, free_steps, dt,
++ method="ep_relax")
++ output = states[-1]
++ _, top5_indices = jax.lax.top_k(output, 5)
++ return {
++ "loss": jnp.mean(jnp.square(output - labels_onehot)),
++ "accuracy": 100.0 * jnp.mean(
++ jnp.argmax(output, -1) == labels, dtype=jnp.float32),
++ "top5accuracy": 100.0 * jnp.mean(
++ jnp.any(top5_indices == labels[:, None], axis=1),
++ dtype=jnp.float32),
++ }
++
++
++def eval_ep_model(state, params, dataset, batch_size, num_classes, free_steps,
++ dt):
++ """Evaluate classification from the settled free-phase EP output."""
++ t0 = time.time()
++ size = len(dataset["image"])
++ steps = size // batch_size
++ indices = jnp.arange(steps * batch_size).reshape((steps, batch_size))
++ metrics = []
++ for index in indices:
++ labels = dataset["label"][index]
++ one_hot = jax.nn.one_hot(labels, num_classes=num_classes)
++ metrics.append(eval_step_ep(
++ state, params, dataset["image"][index], one_hot, labels,
++ free_steps, dt))
++ host = jax.device_get(metrics)
++ summary = {
++ key: float(np.mean([record[key] for record in host]))
++ for key in host[0]
++ }
++ runtime = time.time() - t0
++ nan_diagnostics = [
++ jnp.asarray(jnp.nan, dtype=jnp.float32) for _ in range(len(params))
++ ]
++ return (
++ summary["loss"], summary["accuracy"], summary["top5accuracy"],
++ runtime, nan_diagnostics, nan_diagnostics, nan_diagnostics,
++ nan_diagnostics)
++
+ def ref_grad_and_angle(state, grads, image, labels_onehot, metrics):
+ # Compute a reference backprop gradient (but don't use it).
+ def loss_fn_ref(params, state):
+diff --git a/tests/local_rules_smoke.py b/tests/local_rules_smoke.py
+index 8ce1e02..51d4c9f 100644
+--- a/tests/local_rules_smoke.py
++++ b/tests/local_rules_smoke.py
+@@ -189,6 +189,39 @@ def main():
+ assert ff_scores.shape == (x.shape[0], 3)
+ assert bool(jnp.all(jnp.isfinite(ff_scores)))
+
++ ep_free = model.apply(
++ {"params": params}, x, one_hot, 0.0, 2, 0.5,
++ method="ep_relax")
++ ep_nudged = model.apply(
++ {"params": params}, x, one_hot, 0.5, 1, 0.5, ep_free,
++ method="ep_relax")
++ assert all(
++ bool(jnp.all((state >= 0.0) & (state <= 1.0)))
++ for state in ep_free[1:] + ep_nudged[1:]
++ )
++ ep_grads = jax.grad(
++ lambda candidate: model.apply(
++ {"params": candidate}, ep_free, ep_nudged, 0.5,
++ method="ep_contrastive_objective")
++ )(params)
++ free_readout_input = ep_free[-2].reshape((x.shape[0], -1))
++ nudged_readout_input = ep_nudged[-2].reshape((x.shape[0], -1))
++ ep_manual_kernel = (
++ free_readout_input.T @ ep_free[-1]
++ - nudged_readout_input.T @ ep_nudged[-1]
++ ) / (0.5 * x.shape[0])
++ ep_manual_bias = (ep_free[-1] - ep_nudged[-1]).mean(axis=0) / 0.5
++ ep_readout_leaves = jax.tree_util.tree_leaves(ep_grads["d00"])
++ ep_readout_kernel = next(
++ value for value in ep_readout_leaves if value.ndim == 2)
++ ep_readout_bias = next(
++ value for value in ep_readout_leaves if value.ndim == 1)
++ ep_contrastive_error = max(
++ float(jnp.max(jnp.abs(ep_readout_kernel - ep_manual_kernel))),
++ float(jnp.max(jnp.abs(ep_readout_bias - ep_manual_bias))),
++ )
++ assert ep_contrastive_error < 2e-12, ep_contrastive_error
++
+ report = {
+ "symmetric_fa_bp_relative_error": symmetric_error,
+ "independent_feedback_forward_cosine": feedback_cosine,
+@@ -199,6 +232,7 @@ def main():
+ "pepita_projection_shape": list(pepita_projection.shape),
+ "ff_local_objective_error": ff_objective_error,
+ "ff_score_shape": list(ff_scores.shape),
++ "ep_contrastive_readout_max_error": ep_contrastive_error,
+ "status": "passed",
+ }
+ print(json.dumps(report, indent=2, sort_keys=True))
+diff --git a/train.py b/train.py
+index d3f0de6..08f42c9 100644
+--- a/train.py
++++ b/train.py
+@@ -8,7 +8,7 @@ from absl import logging # for logging
+ from matplotlib.colors import LogNorm
+
+ # Training utils
+-from src import create_train_state, create_local_feedback, train_epoch, eval_model, heatmap_grads_batches, heatmap_grads_epochs, plot_L_or_gamma
++from src import create_train_state, create_local_feedback, train_epoch, eval_model, eval_ep_model, heatmap_grads_batches, heatmap_grads_epochs, plot_L_or_gamma
+
+ # Import configurations
+ # import config # Use this for the old method
+@@ -83,6 +83,8 @@ for experiment_index, seed in enumerate(config.seeds):
+ state, config.train_ds, config.batch_size, input_rng,
+ augmentation_on, config.learning_algorithm, config.num_classes,
+ local_feedback=local_feedback,
++ ep_beta=config.ep_beta, ep_free_steps=config.ep_free_steps,
++ ep_nudge_steps=config.ep_nudge_steps, ep_dt=config.ep_dt,
+ gradient_diagnostics=(config.gradient_diagnostics == "full"))
+ loginfo_and_print('train: \tloss: %.4f, \taccuracy: %.4f, \truntime: %.4f' % (epoch_metrics["loss"], epoch_metrics["accuracy"], train_time))
+ hist['train_loss'][epoch-1], hist['train_accuracy'][epoch-1], hist['train_time'][epoch-1] = epoch_metrics["loss"], epoch_metrics["accuracy"], train_time
+@@ -95,10 +97,15 @@ for experiment_index, seed in enumerate(config.seeds):
+
+ # Evaluate on the validation set after each training epoch
+ rng, input_rng = jax.random.split(rng)
+- val_loss, val_accuracy, val_top5_accuracy, val_time, L10, L20, gamma10, gamma20 = eval_model(
+- state, state.params, config.val_ds, config.batch_size,
+- config.num_classes, input_rng,
+- spectral_diagnostics=(config.spectral_diagnostics == "full"))
++ if config.learning_algorithm == "ep":
++ val_loss, val_accuracy, val_top5_accuracy, val_time, L10, L20, gamma10, gamma20 = eval_ep_model(
++ state, state.params, config.val_ds, config.batch_size,
++ config.num_classes, config.ep_free_steps, config.ep_dt)
++ else:
++ val_loss, val_accuracy, val_top5_accuracy, val_time, L10, L20, gamma10, gamma20 = eval_model(
++ state, state.params, config.val_ds, config.batch_size,
++ config.num_classes, input_rng,
++ spectral_diagnostics=(config.spectral_diagnostics == "full"))
+ loginfo_and_print('val: \tloss: %.4f, \taccuracy: %.4f, \ttop5_accuracy: %.4f, \truntime: %.4f' % (val_loss, val_accuracy, val_top5_accuracy, val_time))
+ loginfo_and_print(f"L20: {[np.round(Li.item(), decimals=4) for Li in L20]}")
+ loginfo_and_print(f"gamma20: {[np.round(gi.item(), decimals=4) for gi in gamma20]}")
+@@ -123,10 +130,15 @@ for experiment_index, seed in enumerate(config.seeds):
+ loginfo_and_print(f"\n====Loading model with best validation accuracy (epoch {best_epoch})====")
+ best_state = checkpoints.restore_checkpoint(ckpt_dir=CKPT_DIR, target=state)
+ rng, input_rng = jax.random.split(rng)
+- test_loss, test_accuracy, test_top5accuracy, test_time, _, _, _, _ = eval_model(
+- best_state, best_state.params, config.test_ds, config.batch_size,
+- config.num_classes, input_rng,
+- spectral_diagnostics=(config.spectral_diagnostics == "full"))
++ if config.learning_algorithm == "ep":
++ test_loss, test_accuracy, test_top5accuracy, test_time, _, _, _, _ = eval_ep_model(
++ best_state, best_state.params, config.test_ds, config.batch_size,
++ config.num_classes, config.ep_free_steps, config.ep_dt)
++ else:
++ test_loss, test_accuracy, test_top5accuracy, test_time, _, _, _, _ = eval_model(
++ best_state, best_state.params, config.test_ds, config.batch_size,
++ config.num_classes, input_rng,
++ spectral_diagnostics=(config.spectral_diagnostics == "full"))
+ hist['test_loss'], hist['test_accuracy'], hist['test_top5accuracy'], hist['test_time'] = test_loss, test_accuracy, test_top5accuracy, test_time
+ loginfo_and_print('test: \tloss: %.4f, \taccuracy: %.4f, \ttop5_accuracy: %.4f, \truntime: %.4f' % (test_loss, test_accuracy, test_top5accuracy, test_time))
+
+--
+2.54.0
+
diff --git a/external/dualprop_patches/0007-crossover-add-reciprocal-clean-KP-training.patch b/external/dualprop_patches/0007-crossover-add-reciprocal-clean-KP-training.patch
new file mode 100644
index 0000000..c869770
--- /dev/null
+++ b/external/dualprop_patches/0007-crossover-add-reciprocal-clean-KP-training.patch
@@ -0,0 +1,321 @@
+From 8b8dfd0fd0a0ba01e66bfca454ca521f68910bd5 Mon Sep 17 00:00:00 2001
+From: YurenHao0426 <Blackhao0426@gmail.com>
+Date: Mon, 27 Jul 2026 13:19:37 -0500
+Subject: [PATCH 07/19] crossover: add reciprocal clean KP training
+
+---
+ config/cli_config.py | 4 +-
+ src/__init__.py | 2 +-
+ src/models.py | 33 ++++++++++++++++
+ src/training_utils.py | 79 ++++++++++++++++++++++++++++++++++++++
+ tests/local_rules_smoke.py | 37 ++++++++++++++++++
+ train.py | 39 +++++++++++++++----
+ 6 files changed, 183 insertions(+), 11 deletions(-)
+
+diff --git a/config/cli_config.py b/config/cli_config.py
+index 795e6a3..56fdb9f 100644
+--- a/config/cli_config.py
++++ b/config/cli_config.py
+@@ -42,7 +42,7 @@ parser.add_argument('--experiment-name', default='test', help='A string denoting
+
+ parser.add_argument('--model', default='VGG16', choices=['VGG16', 'VGGlike', 'CNN', 'miniCNN', 'MLP'], help='')
+
+-parser.add_argument('--learning-algorithm', default='dualprop-lagr-ff', choices=['backprop', 'fa', 'dfa', 'pepita', 'ff', 'ep', 'dualprop-lagr-ff', 'dualprop-raovr-ff', 'dualprop-raovr-dampened-ff'])
++parser.add_argument('--learning-algorithm', default='dualprop-lagr-ff', choices=['backprop', 'fa', 'dfa', 'pepita', 'ff', 'ep', 'clean-kp', 'dualprop-lagr-ff', 'dualprop-raovr-ff', 'dualprop-raovr-dampened-ff'])
+
+ parser.add_argument(
+ '--feedback-seed', default=1729, type=int,
+@@ -156,7 +156,7 @@ elif config.model == "MLP":
+ dense_features = [1024, 1024, config.num_classes]
+
+ # Load model
+-modeltype = {"backprop":cnn_abstract, "fa": cnn_abstract, "dfa": cnn_abstract, "pepita": cnn_abstract, "ff": cnn_abstract, "ep": cnn_abstract, "dualprop-lagr-ff": cnn_dualprop_Lagr_ff, "dualprop-raovr-ff": cnn_dualprop_RAOVR_ff, "dualprop-raovr-dampened-ff": cnn_dualprop_RAOVR_dampened_ff}
++modeltype = {"backprop":cnn_abstract, "fa": cnn_abstract, "dfa": cnn_abstract, "pepita": cnn_abstract, "ff": cnn_abstract, "ep": cnn_abstract, "clean-kp": cnn_abstract, "dualprop-lagr-ff": cnn_dualprop_Lagr_ff, "dualprop-raovr-ff": cnn_dualprop_RAOVR_ff, "dualprop-raovr-dampened-ff": cnn_dualprop_RAOVR_dampened_ff}
+ activation={"relu": relu, "hs": hs, "sigmoid": sigmoid, "tanh": tanh}
+ config.model = modeltype[config.learning_algorithm](loss_func, Conv, Dense, activation[config.activation], config.num_classes, config.beta, config.alpha, config.dtype, config.param_dtype,
+ kernels=kernels, strides=strides, features=features, mp = mp,
+diff --git a/src/__init__.py b/src/__init__.py
+index 58c9adc..7ee7740 100644
+--- a/src/__init__.py
++++ b/src/__init__.py
+@@ -1,2 +1,2 @@
+ from .models import cnn_dualprop_Lagr_ff, cnn_dualprop_RAOVR_ff, cnn_dualprop_RAOVR_dampened_ff, cnn_abstract
+-from .training_utils import create_train_state, create_ff_train_state, create_local_feedback, train_epoch, train_ff_epoch, eval_model, eval_ep_model, eval_ff_model, get_mnist, get_svhn, get_fashionmnist, get_cifar10, get_cifar100, get_imagenet_32x32, heatmap_grads_batches, heatmap_grads_epochs, plot_L_or_gamma
++from .training_utils import create_train_state, create_ff_train_state, create_local_feedback, train_epoch, train_ff_epoch, train_kp_epoch, eval_model, eval_ep_model, eval_ff_model, get_mnist, get_svhn, get_fashionmnist, get_cifar10, get_cifar100, get_imagenet_32x32, heatmap_grads_batches, heatmap_grads_epochs, plot_L_or_gamma
+diff --git a/src/models.py b/src/models.py
+index 68e9804..36b2643 100644
+--- a/src/models.py
++++ b/src/models.py
+@@ -195,6 +195,39 @@ class cnn_abstract(nn.Module, ABC):
+ fields[i] = linear_pullback(linear_field)[0]
+ return fields
+
++ def kp_reciprocal_objective(self, s, forward_linear, teaching_fields):
++ """Recompute KP feedback correlations from local activities only.
++
++ This method is evaluated with Q parameters, but activation gates and
++ max-pool switches come from the cached W forward pass. Its derivative
++ therefore equals the corresponding W local correlation for any values
++ of W and Q, without reading either W or its update.
++ """
++ objective = 0.0
++ batch_size = s[0].shape[0]
++ for i, layer in enumerate(self.layers):
++ child_field = jax.lax.stop_gradient(teaching_fields[i + 1])
++ if i == self.num_layers - 1:
++ linear_field = child_field
++ elif i < self.num_convlayers:
++ _, post_pullback = jax.vjp(
++ lambda z: self.act(self.layers[i].pooling(z)),
++ jax.lax.stop_gradient(forward_linear[i]))
++ linear_field = post_pullback(child_field)[0]
++ else:
++ _, post_pullback = jax.vjp(
++ self.act, jax.lax.stop_gradient(forward_linear[i]))
++ linear_field = post_pullback(child_field)[0]
++
++ x = jax.lax.stop_gradient(s[i])
++ if i < self.num_convlayers:
++ prediction = layer.call_without_pooling(x)
++ else:
++ prediction = layer(x)
++ objective += jnp.sum(
++ prediction * jax.lax.stop_gradient(linear_field))
++ return objective / batch_size
++
+ def pepita_correlation_objective(self, modulated_states, original_linear,
+ modulated_linear, one_hot):
+ """Architecture-compatible PEPITA/ERIN two-presentation update.
+diff --git a/src/training_utils.py b/src/training_utils.py
+index f767b80..e573bb5 100644
+--- a/src/training_utils.py
++++ b/src/training_utils.py
+@@ -347,6 +347,35 @@ def train_epoch(state, train_ds, batch_size, rng, augmentation_on,
+ return state, epoch_metrics_np, runtime
+
+
++def train_kp_epoch(state, feedback_state, train_ds, batch_size, rng,
++ augmentation_on, num_classes, gradient_diagnostics=True):
++ """Train W and reciprocal Q once over a shared minibatch order."""
++ t0 = time.time()
++ size = len(train_ds["image"])
++ steps = size // batch_size
++ perms = jax.random.permutation(rng, size)[:steps * batch_size]
++ perms = perms.reshape((steps, batch_size))
++ metrics = []
++ for permutation in perms:
++ image = train_ds["image"][permutation]
++ labels = train_ds["label"][permutation]
++ one_hot = jax.nn.one_hot(labels, num_classes=num_classes)
++ rng, inf_rng, batch_rng = jax.random.split(rng, 3)
++ per_example_rng = jax.random.split(batch_rng, image.shape[0])
++ state, feedback_state, batch_metrics = train_step_kp(
++ state, feedback_state, image, one_hot, labels, per_example_rng,
++ inf_rng, augmentation_on, gradient_diagnostics)
++ metrics.append(batch_metrics)
++ host = jax.device_get(metrics)
++ summary = {}
++ for key in host[0]:
++ if key == "cosine_sim":
++ summary[key] = [record[key] for record in host]
++ else:
++ summary[key] = np.mean([record[key] for record in host], axis=0)
++ return state, feedback_state, summary, time.time() - t0
++
++
+ def direct_feedback_fields(states, output_field, direct_feedback):
+ """Apply per-hidden fixed output maps without a layerwise reverse chain."""
+ fields = [jnp.zeros_like(states[0])]
+@@ -357,6 +386,56 @@ def direct_feedback_fields(states, output_field, direct_feedback):
+ return fields
+
+
++@jax.jit
++def train_step_kp(state, feedback_state, image, labels_onehot, labels,
++ batch_rng, inf_rng, augmentation_on,
++ gradient_diagnostics):
++ """One simultaneous modified-KP W/Q step."""
++ image = jax.lax.cond(
++ augmentation_on, vmap_augment_train, no_aug, image, batch_rng)
++ states, linear = state.apply_fn(
++ {"params": state.params}, image, method="ff_with_local_cache")
++ states = tree_map(jax.lax.stop_gradient, states)
++ linear = tree_map(jax.lax.stop_gradient, linear)
++
++ def output_loss(logits):
++ return state.apply_fn(
++ {"params": state.params}, logits, labels_onehot,
++ method="output_loss")
++
++ output_field = jax.lax.stop_gradient(jax.grad(output_loss)(states[-1]))
++ teaching_fields = feedback_state.apply_fn(
++ {"params": feedback_state.params}, states, linear, output_field,
++ method="fa_teaching_fields")
++ teaching_fields = tree_map(jax.lax.stop_gradient, teaching_fields)
++
++ forward_grads = jax.grad(
++ lambda params: state.apply_fn(
++ {"params": params}, states, teaching_fields,
++ method="local_correlation_objective")
++ )(state.params)
++ reciprocal_grads = jax.grad(
++ lambda params: feedback_state.apply_fn(
++ {"params": params}, states, linear, teaching_fields,
++ method="kp_reciprocal_objective")
++ )(feedback_state.params)
++ metrics = compute_metrics(
++ image=image, labels_onehot=labels_onehot, labels=labels, state=state)
++ metrics["feedback_forward_cosine"] = cosine_sim_tree(
++ feedback_state.params, state.params)
++ metrics["reciprocal_gradient_error"] = jnp.linalg.norm(
++ jax.flatten_util.ravel_pytree(forward_grads)[0]
++ - jax.flatten_util.ravel_pytree(reciprocal_grads)[0])
++ metrics = jax.lax.cond(
++ gradient_diagnostics, ref_grad_and_angle, no_ref_grad_and_angle,
++ state, forward_grads, image, labels_onehot, metrics)
++
++ # Form both independent local correlations before changing either path.
++ feedback_state = feedback_state.apply_gradients(grads=reciprocal_grads)
++ state = state.apply_gradients(grads=forward_grads)
++ return state, feedback_state, metrics
++
++
+ @partial(
+ jax.jit,
+ static_argnames=("free_steps", "nudge_steps"),
+diff --git a/tests/local_rules_smoke.py b/tests/local_rules_smoke.py
+index 51d4c9f..f40f0ab 100644
+--- a/tests/local_rules_smoke.py
++++ b/tests/local_rules_smoke.py
+@@ -81,6 +81,40 @@ def main():
+ )(params)
+ symmetric_error = relative_error(symmetric_grads, bp_grads)
+ assert symmetric_error < 2e-12, symmetric_error
++ reciprocal_grads = jax.grad(
++ lambda candidate: model.apply(
++ {"params": candidate}, states, linear, symmetric_fields,
++ method="kp_reciprocal_objective")
++ )(params)
++ reciprocal_error = relative_error(reciprocal_grads, symmetric_grads)
++ assert reciprocal_error < 2e-12, reciprocal_error
++ changed_feedback = replace_layer(params, "c01", -0.5)
++ changed_reciprocal_grads = jax.grad(
++ lambda candidate: model.apply(
++ {"params": candidate}, states, linear, symmetric_fields,
++ method="kp_reciprocal_objective")
++ )(changed_feedback)
++ reciprocal_independence_error = relative_error(
++ changed_reciprocal_grads, reciprocal_grads)
++ assert reciprocal_independence_error == 0.0
++
++ kp_optimizer = optax.chain(
++ optax.add_decayed_weights(1e-2),
++ optax.sgd(0.1, momentum=0.9),
++ )
++ kp_w = params
++ kp_q = params
++ kp_w_state = kp_optimizer.init(kp_w)
++ kp_q_state = kp_optimizer.init(kp_q)
++ for _ in range(2):
++ w_updates, kp_w_state = kp_optimizer.update(
++ symmetric_grads, kp_w_state, kp_w)
++ q_updates, kp_q_state = kp_optimizer.update(
++ reciprocal_grads, kp_q_state, kp_q)
++ kp_w = optax.apply_updates(kp_w, w_updates)
++ kp_q = optax.apply_updates(kp_q, q_updates)
++ kp_symmetric_tracking_error = relative_error(kp_q, kp_w)
++ assert kp_symmetric_tracking_error == 0.0
+
+ feedback = create_local_feedback(
+ jax.random.PRNGKey(1729), model, params, (8, 8, 2), "fa", 3)
+@@ -224,6 +258,9 @@ def main():
+
+ report = {
+ "symmetric_fa_bp_relative_error": symmetric_error,
++ "kp_reciprocal_gradient_relative_error": reciprocal_error,
++ "kp_reciprocal_independence_error": reciprocal_independence_error,
++ "kp_symmetric_two_step_tracking_error": kp_symmetric_tracking_error,
+ "independent_feedback_forward_cosine": feedback_cosine,
+ "feedback_independence_error": feedback_independence_error,
+ "detached_local_boundary_error": local_boundary_error,
+diff --git a/train.py b/train.py
+index 08f42c9..6b2d8d8 100644
+--- a/train.py
++++ b/train.py
+@@ -8,7 +8,7 @@ from absl import logging # for logging
+ from matplotlib.colors import LogNorm
+
+ # Training utils
+-from src import create_train_state, create_local_feedback, train_epoch, eval_model, eval_ep_model, heatmap_grads_batches, heatmap_grads_epochs, plot_L_or_gamma
++from src import create_train_state, create_local_feedback, train_epoch, train_kp_epoch, eval_model, eval_ep_model, heatmap_grads_batches, heatmap_grads_epochs, plot_L_or_gamma
+
+ # Import configurations
+ # import config # Use this for the old method
+@@ -54,6 +54,14 @@ for experiment_index, seed in enumerate(config.seeds):
+ jax.random.PRNGKey(config.feedback_seed), config.model, state.params,
+ config.image_dims, config.learning_algorithm, config.num_classes,
+ pepita_projection_scale=config.pepita_projection_scale)
++ feedback_state = None
++ if config.learning_algorithm == "clean-kp":
++ feedback_state = create_train_state(
++ jax.random.PRNGKey(config.feedback_seed), config.model,
++ config.image_dims, config.learning_rate,
++ config.warmup_learning_rate, config.learning_rate_final,
++ config.momentum, config.weight_decay, config.num_epochs,
++ config.warmup_epochs, config.decay_epochs, steps_per_epoch)
+
+
+ del init_rng # Must not be used anymore.
+@@ -64,6 +72,8 @@ for experiment_index, seed in enumerate(config.seeds):
+ 'test_loss': np.nan, 'test_accuracy': np.nan, 'test_top5accuracy': np.nan, 'test_time': np.nan,
+ 'grad_cos_sim_batches': np.zeros((len(state.params), steps_per_epoch*config.num_epochs)),
+ 'grad_cos_sim_epochs': np.zeros((len(state.params), config.num_epochs)),
++ 'feedback_forward_cosine': np.zeros((len(state.params), config.num_epochs)),
++ 'reciprocal_gradient_error': np.zeros(config.num_epochs),
+ 'L10': np.zeros((len(state.params), config.num_epochs)),
+ 'L20': np.zeros((len(state.params), config.num_epochs)),
+ 'gamma10': np.zeros((len(state.params), config.num_epochs)),
+@@ -79,13 +89,21 @@ for experiment_index, seed in enumerate(config.seeds):
+ # Run an optimization step over a training batch
+ # last augument turns off data augmentation for mnist
+ augmentation_on = (config.dataset!="mnist") and (config.dataset!="fashionmnist")
+- state, epoch_metrics, train_time = train_epoch(
+- state, config.train_ds, config.batch_size, input_rng,
+- augmentation_on, config.learning_algorithm, config.num_classes,
+- local_feedback=local_feedback,
+- ep_beta=config.ep_beta, ep_free_steps=config.ep_free_steps,
+- ep_nudge_steps=config.ep_nudge_steps, ep_dt=config.ep_dt,
+- gradient_diagnostics=(config.gradient_diagnostics == "full"))
++ if config.learning_algorithm == "clean-kp":
++ state, feedback_state, epoch_metrics, train_time = train_kp_epoch(
++ state, feedback_state, config.train_ds, config.batch_size,
++ input_rng, augmentation_on, config.num_classes,
++ gradient_diagnostics=(
++ config.gradient_diagnostics == "full"))
++ else:
++ state, epoch_metrics, train_time = train_epoch(
++ state, config.train_ds, config.batch_size, input_rng,
++ augmentation_on, config.learning_algorithm,
++ config.num_classes, local_feedback=local_feedback,
++ ep_beta=config.ep_beta, ep_free_steps=config.ep_free_steps,
++ ep_nudge_steps=config.ep_nudge_steps, ep_dt=config.ep_dt,
++ gradient_diagnostics=(
++ config.gradient_diagnostics == "full"))
+ loginfo_and_print('train: \tloss: %.4f, \taccuracy: %.4f, \truntime: %.4f' % (epoch_metrics["loss"], epoch_metrics["accuracy"], train_time))
+ hist['train_loss'][epoch-1], hist['train_accuracy'][epoch-1], hist['train_time'][epoch-1] = epoch_metrics["loss"], epoch_metrics["accuracy"], train_time
+
+@@ -93,6 +111,11 @@ for experiment_index, seed in enumerate(config.seeds):
+ grad_cos_sim = np.stack(epoch_metrics["cosine_sim"]).T
+ hist["grad_cos_sim_batches"][:,(epoch-1)*steps_per_epoch:(epoch)*steps_per_epoch] = grad_cos_sim
+ hist["grad_cos_sim_epochs"][:,epoch-1] = grad_cos_sim.mean(axis=1)
++ if config.learning_algorithm == "clean-kp":
++ hist["feedback_forward_cosine"][:, epoch - 1] = (
++ epoch_metrics["feedback_forward_cosine"])
++ hist["reciprocal_gradient_error"][epoch - 1] = (
++ epoch_metrics["reciprocal_gradient_error"])
+
+
+ # Evaluate on the validation set after each training epoch
+--
+2.54.0
+
diff --git a/external/dualprop_patches/0008-crossover-add-dynamic-innovation-on-reciprocal-KP.patch b/external/dualprop_patches/0008-crossover-add-dynamic-innovation-on-reciprocal-KP.patch
new file mode 100644
index 0000000..d1da5b2
--- /dev/null
+++ b/external/dualprop_patches/0008-crossover-add-dynamic-innovation-on-reciprocal-KP.patch
@@ -0,0 +1,430 @@
+From d4b231c502fe158752e9294ece450554f42f7cae Mon Sep 17 00:00:00 2001
+From: YurenHao0426 <Blackhao0426@gmail.com>
+Date: Mon, 27 Jul 2026 13:22:51 -0500
+Subject: [PATCH 08/19] crossover: add dynamic innovation on reciprocal KP
+
+---
+ config/cli_config.py | 16 ++-
+ src/__init__.py | 2 +-
+ src/training_utils.py | 205 +++++++++++++++++++++++++++++++++++++
+ tests/local_rules_smoke.py | 25 +++++
+ train.py | 40 +++++++-
+ 5 files changed, 281 insertions(+), 7 deletions(-)
+
+diff --git a/config/cli_config.py b/config/cli_config.py
+index 56fdb9f..4979b37 100644
+--- a/config/cli_config.py
++++ b/config/cli_config.py
+@@ -42,7 +42,7 @@ parser.add_argument('--experiment-name', default='test', help='A string denoting
+
+ parser.add_argument('--model', default='VGG16', choices=['VGG16', 'VGGlike', 'CNN', 'miniCNN', 'MLP'], help='')
+
+-parser.add_argument('--learning-algorithm', default='dualprop-lagr-ff', choices=['backprop', 'fa', 'dfa', 'pepita', 'ff', 'ep', 'clean-kp', 'dualprop-lagr-ff', 'dualprop-raovr-ff', 'dualprop-raovr-dampened-ff'])
++parser.add_argument('--learning-algorithm', default='dualprop-lagr-ff', choices=['backprop', 'fa', 'dfa', 'pepita', 'ff', 'ep', 'clean-kp', 'sdil', 'dualprop-lagr-ff', 'dualprop-raovr-ff', 'dualprop-raovr-dampened-ff'])
+
+ parser.add_argument(
+ '--feedback-seed', default=1729, type=int,
+@@ -76,6 +76,18 @@ parser.add_argument(
+ '--ep-nudge-steps', default=4, type=int,
+ help='Number of nudged-phase EP relaxation steps.')
+
++parser.add_argument(
++ '--sdil-traffic-ratio', default=4.0, type=float,
++ help='Initialization-calibrated traffic/instruction RMS ratio.')
++
++parser.add_argument(
++ '--sdil-traffic-seed', default=4000, type=int,
++ help='Fixed per-cell soma-predictable traffic seed.')
++
++parser.add_argument(
++ '--sdil-calibration-examples', default=64, type=int,
++ help='Neutral examples for the frozen slow affine predictor fit.')
++
+ parser.add_argument(
+ '--gradient-diagnostics', default='full', choices=['none', 'full'],
+ help=('Compute the exact BP reference gradient and layerwise cosine on '
+@@ -156,7 +168,7 @@ elif config.model == "MLP":
+ dense_features = [1024, 1024, config.num_classes]
+
+ # Load model
+-modeltype = {"backprop":cnn_abstract, "fa": cnn_abstract, "dfa": cnn_abstract, "pepita": cnn_abstract, "ff": cnn_abstract, "ep": cnn_abstract, "clean-kp": cnn_abstract, "dualprop-lagr-ff": cnn_dualprop_Lagr_ff, "dualprop-raovr-ff": cnn_dualprop_RAOVR_ff, "dualprop-raovr-dampened-ff": cnn_dualprop_RAOVR_dampened_ff}
++modeltype = {"backprop":cnn_abstract, "fa": cnn_abstract, "dfa": cnn_abstract, "pepita": cnn_abstract, "ff": cnn_abstract, "ep": cnn_abstract, "clean-kp": cnn_abstract, "sdil": cnn_abstract, "dualprop-lagr-ff": cnn_dualprop_Lagr_ff, "dualprop-raovr-ff": cnn_dualprop_RAOVR_ff, "dualprop-raovr-dampened-ff": cnn_dualprop_RAOVR_dampened_ff}
+ activation={"relu": relu, "hs": hs, "sigmoid": sigmoid, "tanh": tanh}
+ config.model = modeltype[config.learning_algorithm](loss_func, Conv, Dense, activation[config.activation], config.num_classes, config.beta, config.alpha, config.dtype, config.param_dtype,
+ kernels=kernels, strides=strides, features=features, mp = mp,
+diff --git a/src/__init__.py b/src/__init__.py
+index 7ee7740..5acc8cd 100644
+--- a/src/__init__.py
++++ b/src/__init__.py
+@@ -1,2 +1,2 @@
+ from .models import cnn_dualprop_Lagr_ff, cnn_dualprop_RAOVR_ff, cnn_dualprop_RAOVR_dampened_ff, cnn_abstract
+-from .training_utils import create_train_state, create_ff_train_state, create_local_feedback, train_epoch, train_ff_epoch, train_kp_epoch, eval_model, eval_ep_model, eval_ff_model, get_mnist, get_svhn, get_fashionmnist, get_cifar10, get_cifar100, get_imagenet_32x32, heatmap_grads_batches, heatmap_grads_epochs, plot_L_or_gamma
++from .training_utils import create_train_state, create_ff_train_state, create_local_feedback, create_sdil_auxiliary, train_epoch, train_ff_epoch, train_kp_epoch, train_sdil_epoch, eval_model, eval_ep_model, eval_ff_model, get_mnist, get_svhn, get_fashionmnist, get_cifar10, get_cifar100, get_imagenet_32x32, heatmap_grads_batches, heatmap_grads_epochs, plot_L_or_gamma
+diff --git a/src/training_utils.py b/src/training_utils.py
+index e573bb5..ac05857 100644
+--- a/src/training_utils.py
++++ b/src/training_utils.py
+@@ -376,6 +376,159 @@ def train_kp_epoch(state, feedback_state, train_ds, batch_size, rng,
+ return state, feedback_state, summary, time.time() - t0
+
+
++def create_sdil_auxiliary(rng, state, feedback_state, image, labels,
++ num_classes, traffic_ratio):
++ """Initialization-only traffic calibration and neutral predictor fit."""
++ one_hot = jax.nn.one_hot(labels, num_classes=num_classes)
++ states, linear = state.apply_fn(
++ {"params": state.params}, image, method="ff_with_local_cache")
++
++ def output_loss(logits):
++ return state.apply_fn(
++ {"params": state.params}, logits, one_hot,
++ method="output_loss")
++
++ output_field = jax.grad(output_loss)(states[-1])
++ instruction = feedback_state.apply_fn(
++ {"params": feedback_state.params}, states, linear, output_field,
++ method="fa_teaching_fields")
++ keys = jax.random.split(rng, len(states) - 2)
++ coefficients = []
++ gains = []
++ slopes = []
++ biases = []
++ realized = []
++ predictor_residual_ratios = []
++ for key, hidden, signal in zip(keys, states[1:-1], instruction[1:-1]):
++ coefficient = jnp.exp(
++ 0.25 * jax.random.normal(
++ key, hidden.shape[1:], dtype=hidden.dtype))
++ signal_rms = jnp.sqrt(jnp.mean(jnp.square(signal)))
++ unscaled_rms = jnp.sqrt(jnp.mean(jnp.square(coefficient * hidden)))
++ gain = traffic_ratio * signal_rms / jnp.maximum(unscaled_rms, 1e-30)
++ traffic = gain * coefficient * hidden
++ hidden_mean = jnp.mean(hidden, axis=0)
++ traffic_mean = jnp.mean(traffic, axis=0)
++ centered_hidden = hidden - hidden_mean
++ centered_traffic = traffic - traffic_mean
++ variance = jnp.mean(jnp.square(centered_hidden), axis=0)
++ covariance = jnp.mean(
++ centered_hidden * centered_traffic, axis=0)
++ slope = jnp.where(
++ variance > 1e-12,
++ covariance / jnp.maximum(variance, 1e-12),
++ jnp.zeros_like(variance))
++ bias = traffic_mean - slope * hidden_mean
++ residual = traffic - (slope * hidden + bias)
++ coefficients.append(jax.lax.stop_gradient(coefficient))
++ gains.append(jax.lax.stop_gradient(gain))
++ slopes.append(jax.lax.stop_gradient(slope))
++ biases.append(jax.lax.stop_gradient(bias))
++ realized.append(jnp.sqrt(jnp.mean(jnp.square(traffic)))
++ / jnp.maximum(signal_rms, 1e-30))
++ predictor_residual_ratios.append(
++ jnp.sqrt(jnp.mean(jnp.square(residual)))
++ / jnp.maximum(
++ jnp.sqrt(jnp.mean(jnp.square(traffic))), 1e-30))
++ auxiliary = {
++ "coefficients": tuple(coefficients),
++ "gains": tuple(gains),
++ "slopes": tuple(slopes),
++ "biases": tuple(biases),
++ }
++ report = {
++ "traffic_ratio_target": float(traffic_ratio),
++ "realized_traffic_instruction_rms_ratio": [
++ float(value) for value in jax.device_get(realized)],
++ "predictor_residual_traffic_rms_ratio": [
++ float(value)
++ for value in jax.device_get(predictor_residual_ratios)],
++ "observations": int(image.shape[0]),
++ "instruction_observations_for_predictor": 0,
++ }
++ return auxiliary, report
++
++
++def sdil_innovation_fields(states, instruction, auxiliary):
++ """Paired-neutral affine projection for every hidden population."""
++ fields = [instruction[0]]
++ pre_power = jnp.asarray(0.0, dtype=states[0].dtype)
++ post_power = jnp.asarray(0.0, dtype=states[0].dtype)
++ traffic_power = jnp.asarray(0.0, dtype=states[0].dtype)
++ maximum_post_slope = jnp.asarray(0.0, dtype=states[0].dtype)
++ for hidden, signal, coefficient, gain, slope, bias in zip(
++ states[1:-1], instruction[1:-1],
++ auxiliary["coefficients"], auxiliary["gains"],
++ auxiliary["slopes"], auxiliary["biases"]):
++ traffic = gain * coefficient * hidden
++ neutral = traffic - (slope * hidden + bias)
++ centered_hidden = hidden - jnp.mean(hidden, axis=0)
++ centered_neutral = neutral - jnp.mean(neutral, axis=0)
++ variance = jnp.mean(jnp.square(centered_hidden), axis=0)
++ covariance = jnp.mean(
++ centered_hidden * centered_neutral, axis=0)
++ correction = jnp.where(
++ variance > 1e-12,
++ covariance / jnp.maximum(variance, 1e-12),
++ jnp.zeros_like(variance))
++ remainder = centered_neutral - correction * centered_hidden
++ centered_remainder = remainder - jnp.mean(remainder, axis=0)
++ post_covariance = jnp.mean(
++ centered_hidden * centered_remainder, axis=0)
++ post_slope = jnp.where(
++ variance > 1e-12,
++ post_covariance / jnp.maximum(variance, 1e-12),
++ jnp.zeros_like(variance))
++ maximum_post_slope = jnp.maximum(
++ maximum_post_slope, jnp.max(jnp.abs(post_slope)))
++ pre_power += jnp.sum(jnp.square(neutral))
++ post_power += jnp.sum(jnp.square(remainder))
++ traffic_power += jnp.sum(jnp.square(traffic))
++ fields.append(signal + remainder)
++ fields.append(instruction[-1])
++ report = {
++ "pre_projection_traffic_rms_ratio": jnp.sqrt(
++ pre_power / jnp.maximum(traffic_power, 1e-30)),
++ "post_projection_traffic_rms_ratio": jnp.sqrt(
++ post_power / jnp.maximum(traffic_power, 1e-30)),
++ "max_absolute_post_projection_soma_slope": maximum_post_slope,
++ "instruction_observations": jnp.asarray(0, dtype=jnp.int32),
++ "observations": jnp.asarray(states[0].shape[0], dtype=jnp.int32),
++ }
++ return fields, report
++
++
++def train_sdil_epoch(state, feedback_state, auxiliary, train_ds, batch_size,
++ rng, augmentation_on, num_classes,
++ gradient_diagnostics=True):
++ """Train dynamic innovation W/Q once over a shared minibatch order."""
++ t0 = time.time()
++ size = len(train_ds["image"])
++ steps = size // batch_size
++ perms = jax.random.permutation(rng, size)[:steps * batch_size]
++ perms = perms.reshape((steps, batch_size))
++ metrics = []
++ for permutation in perms:
++ image = train_ds["image"][permutation]
++ labels = train_ds["label"][permutation]
++ one_hot = jax.nn.one_hot(labels, num_classes=num_classes)
++ rng, inf_rng, batch_rng = jax.random.split(rng, 3)
++ per_example_rng = jax.random.split(batch_rng, image.shape[0])
++ state, feedback_state, batch_metrics = train_step_sdil(
++ state, feedback_state, auxiliary, image, one_hot, labels,
++ per_example_rng, inf_rng, augmentation_on,
++ gradient_diagnostics)
++ metrics.append(batch_metrics)
++ host = jax.device_get(metrics)
++ summary = {}
++ for key in host[0]:
++ if key == "cosine_sim":
++ summary[key] = [record[key] for record in host]
++ else:
++ summary[key] = np.mean([record[key] for record in host], axis=0)
++ return state, feedback_state, summary, time.time() - t0
++
++
+ def direct_feedback_fields(states, output_field, direct_feedback):
+ """Apply per-hidden fixed output maps without a layerwise reverse chain."""
+ fields = [jnp.zeros_like(states[0])]
+@@ -436,6 +589,58 @@ def train_step_kp(state, feedback_state, image, labels_onehot, labels,
+ return state, feedback_state, metrics
+
+
++@jax.jit
++def train_step_sdil(state, feedback_state, auxiliary, image, labels_onehot,
++ labels, batch_rng, inf_rng, augmentation_on,
++ gradient_diagnostics):
++ """One simultaneous dynamic-innovation W/Q step."""
++ image = jax.lax.cond(
++ augmentation_on, vmap_augment_train, no_aug, image, batch_rng)
++ states, linear = state.apply_fn(
++ {"params": state.params}, image, method="ff_with_local_cache")
++ states = tree_map(jax.lax.stop_gradient, states)
++ linear = tree_map(jax.lax.stop_gradient, linear)
++
++ def output_loss(logits):
++ return state.apply_fn(
++ {"params": state.params}, logits, labels_onehot,
++ method="output_loss")
++
++ output_field = jax.lax.stop_gradient(jax.grad(output_loss)(states[-1]))
++ instruction = feedback_state.apply_fn(
++ {"params": feedback_state.params}, states, linear, output_field,
++ method="fa_teaching_fields")
++ instruction = tree_map(jax.lax.stop_gradient, instruction)
++ teaching_fields, projection = sdil_innovation_fields(
++ states, instruction, auxiliary)
++ teaching_fields = tree_map(jax.lax.stop_gradient, teaching_fields)
++
++ forward_grads = jax.grad(
++ lambda params: state.apply_fn(
++ {"params": params}, states, teaching_fields,
++ method="local_correlation_objective")
++ )(state.params)
++ reciprocal_grads = jax.grad(
++ lambda params: feedback_state.apply_fn(
++ {"params": params}, states, linear, teaching_fields,
++ method="kp_reciprocal_objective")
++ )(feedback_state.params)
++ metrics = compute_metrics(
++ image=image, labels_onehot=labels_onehot, labels=labels, state=state)
++ metrics.update(projection)
++ metrics["feedback_forward_cosine"] = cosine_sim_tree(
++ feedback_state.params, state.params)
++ metrics["reciprocal_gradient_error"] = jnp.linalg.norm(
++ jax.flatten_util.ravel_pytree(forward_grads)[0]
++ - jax.flatten_util.ravel_pytree(reciprocal_grads)[0])
++ metrics = jax.lax.cond(
++ gradient_diagnostics, ref_grad_and_angle, no_ref_grad_and_angle,
++ state, forward_grads, image, labels_onehot, metrics)
++ feedback_state = feedback_state.apply_gradients(grads=reciprocal_grads)
++ state = state.apply_gradients(grads=forward_grads)
++ return state, feedback_state, metrics
++
++
+ @partial(
+ jax.jit,
+ static_argnames=("free_steps", "nudge_steps"),
+diff --git a/tests/local_rules_smoke.py b/tests/local_rules_smoke.py
+index f40f0ab..3763ee8 100644
+--- a/tests/local_rules_smoke.py
++++ b/tests/local_rules_smoke.py
+@@ -19,6 +19,7 @@ from src.training_utils import (
+ create_local_feedback,
+ direct_feedback_fields,
+ ff_overlay,
++ sdil_innovation_fields,
+ )
+
+
+@@ -116,6 +117,28 @@ def main():
+ kp_symmetric_tracking_error = relative_error(kp_q, kp_w)
+ assert kp_symmetric_tracking_error == 0.0
+
++ hidden_states = states[1:-1]
++ sdil_auxiliary = {
++ "coefficients": tuple(jnp.ones(hidden.shape[1:], hidden.dtype)
++ for hidden in hidden_states),
++ "gains": tuple(jnp.asarray(4.0, hidden.dtype)
++ for hidden in hidden_states),
++ # Leave a one-times-soma neutral residual for the fast projection.
++ "slopes": tuple(jnp.full(hidden.shape[1:], 3.0, hidden.dtype)
++ for hidden in hidden_states),
++ "biases": tuple(jnp.full(hidden.shape[1:], 0.2, hidden.dtype)
++ for hidden in hidden_states),
++ }
++ sdil_fields, sdil_projection = sdil_innovation_fields(
++ states, symmetric_fields, sdil_auxiliary)
++ sdil_projection_identity_error = relative_error(
++ sdil_fields[1:-1], symmetric_fields[1:-1])
++ assert sdil_projection_identity_error < 2e-12
++ assert float(
++ sdil_projection["max_absolute_post_projection_soma_slope"]
++ ) < 2e-12
++ assert int(sdil_projection["instruction_observations"]) == 0
++
+ feedback = create_local_feedback(
+ jax.random.PRNGKey(1729), model, params, (8, 8, 2), "fa", 3)
+ feedback_cosine = float(
+@@ -261,6 +284,8 @@ def main():
+ "kp_reciprocal_gradient_relative_error": reciprocal_error,
+ "kp_reciprocal_independence_error": reciprocal_independence_error,
+ "kp_symmetric_two_step_tracking_error": kp_symmetric_tracking_error,
++ "sdil_projection_instruction_identity_error": (
++ sdil_projection_identity_error),
+ "independent_feedback_forward_cosine": feedback_cosine,
+ "feedback_independence_error": feedback_independence_error,
+ "detached_local_boundary_error": local_boundary_error,
+diff --git a/train.py b/train.py
+index 6b2d8d8..b1df1b0 100644
+--- a/train.py
++++ b/train.py
+@@ -8,7 +8,7 @@ from absl import logging # for logging
+ from matplotlib.colors import LogNorm
+
+ # Training utils
+-from src import create_train_state, create_local_feedback, train_epoch, train_kp_epoch, eval_model, eval_ep_model, heatmap_grads_batches, heatmap_grads_epochs, plot_L_or_gamma
++from src import create_train_state, create_local_feedback, create_sdil_auxiliary, train_epoch, train_kp_epoch, train_sdil_epoch, eval_model, eval_ep_model, heatmap_grads_batches, heatmap_grads_epochs, plot_L_or_gamma
+
+ # Import configurations
+ # import config # Use this for the old method
+@@ -55,13 +55,26 @@ for experiment_index, seed in enumerate(config.seeds):
+ config.image_dims, config.learning_algorithm, config.num_classes,
+ pepita_projection_scale=config.pepita_projection_scale)
+ feedback_state = None
+- if config.learning_algorithm == "clean-kp":
++ if config.learning_algorithm in ("clean-kp", "sdil"):
+ feedback_state = create_train_state(
+ jax.random.PRNGKey(config.feedback_seed), config.model,
+ config.image_dims, config.learning_rate,
+ config.warmup_learning_rate, config.learning_rate_final,
+ config.momentum, config.weight_decay, config.num_epochs,
+ config.warmup_epochs, config.decay_epochs, steps_per_epoch)
++ sdil_auxiliary = None
++ sdil_initialization = None
++ if config.learning_algorithm == "sdil":
++ count = config.sdil_calibration_examples
++ if count < 2 or count > len(config.train_ds["image"]):
++ raise ValueError("invalid --sdil-calibration-examples")
++ if config.sdil_traffic_ratio <= 0:
++ raise ValueError("--sdil-traffic-ratio must be positive")
++ sdil_auxiliary, sdil_initialization = create_sdil_auxiliary(
++ jax.random.PRNGKey(config.sdil_traffic_seed), state,
++ feedback_state, config.train_ds["image"][:count],
++ config.train_ds["label"][:count], config.num_classes,
++ config.sdil_traffic_ratio)
+
+
+ del init_rng # Must not be used anymore.
+@@ -74,10 +87,14 @@ for experiment_index, seed in enumerate(config.seeds):
+ 'grad_cos_sim_epochs': np.zeros((len(state.params), config.num_epochs)),
+ 'feedback_forward_cosine': np.zeros((len(state.params), config.num_epochs)),
+ 'reciprocal_gradient_error': np.zeros(config.num_epochs),
++ 'pre_projection_traffic_rms_ratio': np.zeros(config.num_epochs),
++ 'post_projection_traffic_rms_ratio': np.zeros(config.num_epochs),
++ 'max_absolute_post_projection_soma_slope': np.zeros(config.num_epochs),
+ 'L10': np.zeros((len(state.params), config.num_epochs)),
+ 'L20': np.zeros((len(state.params), config.num_epochs)),
+ 'gamma10': np.zeros((len(state.params), config.num_epochs)),
+- 'gamma20': np.zeros((len(state.params), config.num_epochs))}
++ 'gamma20': np.zeros((len(state.params), config.num_epochs)),
++ 'sdil_initialization': sdil_initialization}
+
+ best_accuracy, best_epoch = 0, 0
+ epoch = 0
+@@ -95,6 +112,13 @@ for experiment_index, seed in enumerate(config.seeds):
+ input_rng, augmentation_on, config.num_classes,
+ gradient_diagnostics=(
+ config.gradient_diagnostics == "full"))
++ elif config.learning_algorithm == "sdil":
++ state, feedback_state, epoch_metrics, train_time = train_sdil_epoch(
++ state, feedback_state, sdil_auxiliary, config.train_ds,
++ config.batch_size, input_rng, augmentation_on,
++ config.num_classes,
++ gradient_diagnostics=(
++ config.gradient_diagnostics == "full"))
+ else:
+ state, epoch_metrics, train_time = train_epoch(
+ state, config.train_ds, config.batch_size, input_rng,
+@@ -111,11 +135,19 @@ for experiment_index, seed in enumerate(config.seeds):
+ grad_cos_sim = np.stack(epoch_metrics["cosine_sim"]).T
+ hist["grad_cos_sim_batches"][:,(epoch-1)*steps_per_epoch:(epoch)*steps_per_epoch] = grad_cos_sim
+ hist["grad_cos_sim_epochs"][:,epoch-1] = grad_cos_sim.mean(axis=1)
+- if config.learning_algorithm == "clean-kp":
++ if config.learning_algorithm in ("clean-kp", "sdil"):
+ hist["feedback_forward_cosine"][:, epoch - 1] = (
+ epoch_metrics["feedback_forward_cosine"])
+ hist["reciprocal_gradient_error"][epoch - 1] = (
+ epoch_metrics["reciprocal_gradient_error"])
++ if config.learning_algorithm == "sdil":
++ hist["pre_projection_traffic_rms_ratio"][epoch - 1] = (
++ epoch_metrics["pre_projection_traffic_rms_ratio"])
++ hist["post_projection_traffic_rms_ratio"][epoch - 1] = (
++ epoch_metrics["post_projection_traffic_rms_ratio"])
++ hist["max_absolute_post_projection_soma_slope"][epoch - 1] = (
++ epoch_metrics[
++ "max_absolute_post_projection_soma_slope"])
+
+
+ # Evaluate on the validation set after each training epoch
+--
+2.54.0
+
diff --git a/external/dualprop_patches/0009-crossover-isolate-formal-validation-from-test.patch b/external/dualprop_patches/0009-crossover-isolate-formal-validation-from-test.patch
new file mode 100644
index 0000000..03a826a
--- /dev/null
+++ b/external/dualprop_patches/0009-crossover-isolate-formal-validation-from-test.patch
@@ -0,0 +1,135 @@
+From 9fc597c20d6b758ce6b0c3a326b8bc1b9262b0d0 Mon Sep 17 00:00:00 2001
+From: YurenHao0426 <Blackhao0426@gmail.com>
+Date: Mon, 27 Jul 2026 13:25:10 -0500
+Subject: [PATCH 09/19] crossover: isolate formal validation from test
+
+---
+ config/cli_config.py | 10 ++++++++++
+ train.py | 43 ++++++++++++++++++++++++++-----------------
+ train_ff.py | 22 +++++++++++++++-------
+ 3 files changed, 51 insertions(+), 24 deletions(-)
+
+diff --git a/config/cli_config.py b/config/cli_config.py
+index 4979b37..67a04fe 100644
+--- a/config/cli_config.py
++++ b/config/cli_config.py
+@@ -100,6 +100,16 @@ parser.add_argument(
+ 'validation evaluation. The author-compatible default is full; '
+ 'use none for diagnostic-free timed crossover runs.'))
+
++parser.add_argument(
++ '--test-policy', default='author', choices=['author', 'none'],
++ help=('author restores the best-validation checkpoint and evaluates test; '
++ 'none keeps test completely untouched for formal validation runs.'))
++
++parser.add_argument(
++ '--early-stop-policy', default='author', choices=['author', 'none'],
++ help=('author retains the original severe-validation-drop stop; none runs '
++ 'the full schedule unless a nonfinite training loss occurs.'))
++
+ dtypes = {'bfloat16': jnp.bfloat16, 'float16':jnp.float16, 'float32':jnp.float32}
+ parser.add_argument('--dtype', default='float32', choices=dtypes.keys())
+ parser.add_argument('--param-dtype', default='float32', choices=['bfloat16', 'float16', 'float32'])
+diff --git a/train.py b/train.py
+index b1df1b0..0b23304 100644
+--- a/train.py
++++ b/train.py
+@@ -43,7 +43,7 @@ for experiment_index, seed in enumerate(config.seeds):
+ print(msg)
+
+ loginfo_and_print(f"\n\tStarting experiment {experiment_index+1}/{len(config.seeds)}. Current seed is {seed}")
+- loginfo_and_print(f"\model settings:\n{config.model}")
++ loginfo_and_print(f"\nmodel settings:\n{config.model}")
+ rng = jax.random.PRNGKey(seed)
+ rng, init_rng = jax.random.split(rng)
+ # Define model
+@@ -179,23 +179,32 @@ for experiment_index, seed in enumerate(config.seeds):
+ loginfo_and_print("NaN or Inf encountered. Terminating training loop early")
+ break
+ if (epoch>5) and (hist["val_accuracy"][epoch-1] < 0.5*best_accuracy):
+- loginfo_and_print("Terminating training early as validation accuracy has drastically dropped")
+- break
+-
+- loginfo_and_print(f"\n====Loading model with best validation accuracy (epoch {best_epoch})====")
+- best_state = checkpoints.restore_checkpoint(ckpt_dir=CKPT_DIR, target=state)
+- rng, input_rng = jax.random.split(rng)
+- if config.learning_algorithm == "ep":
+- test_loss, test_accuracy, test_top5accuracy, test_time, _, _, _, _ = eval_ep_model(
+- best_state, best_state.params, config.test_ds, config.batch_size,
+- config.num_classes, config.ep_free_steps, config.ep_dt)
++ if config.early_stop_policy == "author":
++ loginfo_and_print("Terminating training early as validation accuracy has drastically dropped")
++ break
++
++ hist["best_validation_accuracy"] = best_accuracy
++ hist["best_epoch"] = best_epoch
++ hist["epochs_completed"] = epoch
++ if config.test_policy == "author":
++ loginfo_and_print(f"\n====Loading model with best validation accuracy (epoch {best_epoch})====")
++ best_state = checkpoints.restore_checkpoint(ckpt_dir=CKPT_DIR, target=state)
++ rng, input_rng = jax.random.split(rng)
++ if config.learning_algorithm == "ep":
++ test_loss, test_accuracy, test_top5accuracy, test_time, _, _, _, _ = eval_ep_model(
++ best_state, best_state.params, config.test_ds,
++ config.batch_size, config.num_classes, config.ep_free_steps,
++ config.ep_dt)
++ else:
++ test_loss, test_accuracy, test_top5accuracy, test_time, _, _, _, _ = eval_model(
++ best_state, best_state.params, config.test_ds,
++ config.batch_size, config.num_classes, input_rng,
++ spectral_diagnostics=(
++ config.spectral_diagnostics == "full"))
++ hist['test_loss'], hist['test_accuracy'], hist['test_top5accuracy'], hist['test_time'] = test_loss, test_accuracy, test_top5accuracy, test_time
++ loginfo_and_print('test: \tloss: %.4f, \taccuracy: %.4f, \ttop5_accuracy: %.4f, \truntime: %.4f' % (test_loss, test_accuracy, test_top5accuracy, test_time))
+ else:
+- test_loss, test_accuracy, test_top5accuracy, test_time, _, _, _, _ = eval_model(
+- best_state, best_state.params, config.test_ds, config.batch_size,
+- config.num_classes, input_rng,
+- spectral_diagnostics=(config.spectral_diagnostics == "full"))
+- hist['test_loss'], hist['test_accuracy'], hist['test_top5accuracy'], hist['test_time'] = test_loss, test_accuracy, test_top5accuracy, test_time
+- loginfo_and_print('test: \tloss: %.4f, \taccuracy: %.4f, \ttop5_accuracy: %.4f, \truntime: %.4f' % (test_loss, test_accuracy, test_top5accuracy, test_time))
++ loginfo_and_print("\n====Test evaluation disabled by validation-only policy====")
+
+ if (config.learning_algorithm != "backprop"
+ and config.gradient_diagnostics == "full"):
+diff --git a/train_ff.py b/train_ff.py
+index 3839dd3..ada45a3 100644
+--- a/train_ff.py
++++ b/train_ff.py
+@@ -75,9 +75,12 @@ for experiment_index, seed in enumerate(config.seeds):
+ validation_accuracy, validation_time = eval_ff_model(
+ state, config.val_ds, config.batch_size, config.num_classes,
+ config.ff_score_from_layer)
+- test_accuracy, test_time = eval_ff_model(
+- state, config.test_ds, config.batch_size, config.num_classes,
+- config.ff_score_from_layer)
++ test_accuracy = np.nan
++ test_time = np.nan
++ if config.test_policy == "author":
++ test_accuracy, test_time = eval_ff_model(
++ state, config.test_ds, config.batch_size, config.num_classes,
++ config.ff_score_from_layer)
+ history["final"] = {
+ "validation_accuracy": validation_accuracy,
+ "validation_time": validation_time,
+@@ -85,8 +88,13 @@ for experiment_index, seed in enumerate(config.seeds):
+ "test_time": test_time,
+ "train_and_eval_wall": time.time() - started,
+ }
+- print(
+- f"final val_accuracy={validation_accuracy:.3f}% "
+- f"test_accuracy={test_accuracy:.3f}%",
+- flush=True)
++ if config.test_policy == "author":
++ print(
++ f"final val_accuracy={validation_accuracy:.3f}% "
++ f"test_accuracy={test_accuracy:.3f}%",
++ flush=True)
++ else:
++ print(
++ f"final val_accuracy={validation_accuracy:.3f}% test=disabled",
++ flush=True)
+ np.save(os.path.join(outpath, "hist.npy"), history)
+--
+2.54.0
+
diff --git a/external/dualprop_patches/0010-crossover-add-paper-PEPITA-learning-schedule.patch b/external/dualprop_patches/0010-crossover-add-paper-PEPITA-learning-schedule.patch
new file mode 100644
index 0000000..73f17d1
--- /dev/null
+++ b/external/dualprop_patches/0010-crossover-add-paper-PEPITA-learning-schedule.patch
@@ -0,0 +1,110 @@
+From 17fc860412aa78b16b81f7d2476644b6650b2bd7 Mon Sep 17 00:00:00 2001
+From: YurenHao0426 <Blackhao0426@gmail.com>
+Date: Mon, 27 Jul 2026 13:26:52 -0500
+Subject: [PATCH 10/19] crossover: add paper PEPITA learning schedule
+
+---
+ config/cli_config.py | 8 ++++++++
+ src/training_utils.py | 26 +++++++++++++++++++-------
+ train.py | 10 ++++++++--
+ 3 files changed, 35 insertions(+), 9 deletions(-)
+
+diff --git a/config/cli_config.py b/config/cli_config.py
+index 67a04fe..b38ff7b 100644
+--- a/config/cli_config.py
++++ b/config/cli_config.py
+@@ -24,6 +24,11 @@ parser.add_argument('--decay-epochs', default=None, type=int, help='')
+
+ parser.add_argument('--warmup-epochs', default=0, type=int, help='')
+
++parser.add_argument(
++ '--optimizer-schedule', default='author',
++ choices=['author', 'pepita'],
++ help='Author warmup/cosine schedule or PEPITA epoch-60/90 drops.')
++
+ parser.add_argument('--momentum', default=0.9, type=float, const=None, action='store', nargs='?', help='')
+
+ parser.add_argument('--weight-decay', default=5e-4, type=float, help='')
+@@ -127,6 +132,9 @@ datasets = dict(fashionmnist=get_fashionmnist, mnist=get_mnist, svhn=get_svhn, c
+ parser.add_argument('--dataset', '-d', choices=datasets.keys(), default='cifar10', help='')
+
+ config = parser.parse_args()
++if (config.optimizer_schedule == "pepita"
++ and config.learning_algorithm != "pepita"):
++ parser.error("--optimizer-schedule pepita is restricted to PEPITA")
+
+ # Note for comparisson to previous work we pad MNIST and FMNIST to size (32,32,1)
+ imagedims = {"fashionmnist": (28,28,1), "mnist": (32,32,1), "svhn": (32,32,3), "cifar10": (32,32,3), "cifar100": (32,32,3), "imagenet_32x32": (32,32,3)}
+diff --git a/src/training_utils.py b/src/training_utils.py
+index ac05857..271a81e 100644
+--- a/src/training_utils.py
++++ b/src/training_utils.py
+@@ -197,18 +197,30 @@ def get_imagenet_32x32(dtype, percent_train=95, percent_val=5):
+ test_ds['image'] = (test_ds['image'] - mean_data)/std_data
+ return train_ds, val_ds, test_ds
+
+-def create_train_state(rng, model, image_dims, lr, wlr, lrf, momentum, weight_decay, num_epochs, warmup_epochs, decay_epochs, steps_per_epoch):
++def create_train_state(rng, model, image_dims, lr, wlr, lrf, momentum,
++ weight_decay, num_epochs, warmup_epochs, decay_epochs,
++ steps_per_epoch, schedule_mode="author"):
+ """Creates initial `TrainState`."""
+ w, h, ch = image_dims
+ x = jnp.ones([1, w, h, ch])
+ params = model.init(rng, x)['params']
+
+- schedule = optax.warmup_cosine_decay_schedule(
+- init_value=wlr,
+- peak_value=lr,
+- warmup_steps=warmup_epochs*steps_per_epoch,
+- decay_steps = (decay_epochs)*steps_per_epoch,
+- end_value=lrf)
++ if schedule_mode == "author":
++ schedule = optax.warmup_cosine_decay_schedule(
++ init_value=wlr,
++ peak_value=lr,
++ warmup_steps=warmup_epochs * steps_per_epoch,
++ decay_steps=decay_epochs * steps_per_epoch,
++ end_value=lrf)
++ elif schedule_mode == "pepita":
++ schedule = optax.piecewise_constant_schedule(
++ init_value=lr,
++ boundaries_and_scales={
++ 60 * steps_per_epoch: 0.1,
++ 90 * steps_per_epoch: 0.1,
++ })
++ else:
++ raise ValueError(f"unknown optimizer schedule: {schedule_mode}")
+
+ tx = optax.chain(
+ optax.add_decayed_weights(weight_decay=weight_decay, mask=None),
+diff --git a/train.py b/train.py
+index 0b23304..2a03ada 100644
+--- a/train.py
++++ b/train.py
+@@ -49,7 +49,12 @@ for experiment_index, seed in enumerate(config.seeds):
+ # Define model
+ steps_per_epoch = len(config.train_ds['image']) // config.batch_size
+
+- state = create_train_state(init_rng, config.model, config.image_dims, config.learning_rate, config.warmup_learning_rate, config.learning_rate_final, config.momentum, config.weight_decay, config.num_epochs, config.warmup_epochs, config.decay_epochs, steps_per_epoch)
++ state = create_train_state(
++ init_rng, config.model, config.image_dims, config.learning_rate,
++ config.warmup_learning_rate, config.learning_rate_final,
++ config.momentum, config.weight_decay, config.num_epochs,
++ config.warmup_epochs, config.decay_epochs, steps_per_epoch,
++ schedule_mode=config.optimizer_schedule)
+ local_feedback = create_local_feedback(
+ jax.random.PRNGKey(config.feedback_seed), config.model, state.params,
+ config.image_dims, config.learning_algorithm, config.num_classes,
+@@ -61,7 +66,8 @@ for experiment_index, seed in enumerate(config.seeds):
+ config.image_dims, config.learning_rate,
+ config.warmup_learning_rate, config.learning_rate_final,
+ config.momentum, config.weight_decay, config.num_epochs,
+- config.warmup_epochs, config.decay_epochs, steps_per_epoch)
++ config.warmup_epochs, config.decay_epochs, steps_per_epoch,
++ schedule_mode=config.optimizer_schedule)
+ sdil_auxiliary = None
+ sdil_initialization = None
+ if config.learning_algorithm == "sdil":
+--
+2.54.0
+
diff --git a/external/dualprop_patches/0011-experiment-register-plain-CNN-P1-selector-grid.patch b/external/dualprop_patches/0011-experiment-register-plain-CNN-P1-selector-grid.patch
new file mode 100644
index 0000000..4f79f09
--- /dev/null
+++ b/external/dualprop_patches/0011-experiment-register-plain-CNN-P1-selector-grid.patch
@@ -0,0 +1,247 @@
+From bc3b656233c8731e7b8182087cf74a0079baadbd Mon Sep 17 00:00:00 2001
+From: YurenHao0426 <Blackhao0426@gmail.com>
+Date: Mon, 27 Jul 2026 13:28:00 -0500
+Subject: [PATCH 11/19] experiment: register plain CNN P1 selector grid
+
+---
+ crossover_grid.py | 228 ++++++++++++++++++++++++++++++++++++++++++++++
+ 1 file changed, 228 insertions(+)
+ create mode 100644 crossover_grid.py
+
+diff --git a/crossover_grid.py b/crossover_grid.py
+new file mode 100644
+index 0000000..dae3316
+--- /dev/null
++++ b/crossover_grid.py
+@@ -0,0 +1,228 @@
++#!/usr/bin/env python3
++"""Immutable command registry for the plain-CNN crossover stages."""
++import argparse
++import glob
++import hashlib
++import json
++import os
++import subprocess
++import sys
++import time
++
++
++ROOT = os.path.dirname(os.path.abspath(__file__))
++MAIN_ROOT = "/home/yurenh2/sdil"
++PROTOCOL = os.path.join(MAIN_ROOT, "PLAIN_CNN_CROSSOVER.md")
++
++
++def output(*args):
++ return subprocess.run(
++ args, cwd=ROOT, check=True, capture_output=True,
++ text=True).stdout.strip()
++
++
++def sha256(path):
++ digest = hashlib.sha256()
++ with open(path, "rb") as handle:
++ for chunk in iter(lambda: handle.read(1024 * 1024), b""):
++ digest.update(chunk)
++ return digest.hexdigest()
++
++
++def source_report():
++ commit = output("git", "rev-parse", "HEAD")
++ dirty = output(
++ "git", "status", "--porcelain", "--untracked-files=no")
++ if dirty:
++ raise RuntimeError("crossover runs require clean tracked author source")
++ main_commit = subprocess.run(
++ ["git", "rev-parse", "HEAD"], cwd=MAIN_ROOT, check=True,
++ capture_output=True, text=True).stdout.strip()
++ main_dirty = subprocess.run(
++ ["git", "status", "--porcelain", "--untracked-files=no"],
++ cwd=MAIN_ROOT, check=True, capture_output=True,
++ text=True).stdout.strip()
++ if main_dirty:
++ raise RuntimeError("crossover runs require clean tracked main source")
++ return {
++ "author_crossover_commit": commit,
++ "main_commit": main_commit,
++ "protocol_path": PROTOCOL,
++ "protocol_sha256": sha256(PROTOCOL),
++ "python": sys.executable,
++ "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"),
++ }
++
++
++def rate_tag(rate):
++ return f"{rate:g}".replace(".", "p")
++
++
++def p1_jobs():
++ # Expensive jobs are deliberately interleaved across modulo shards.
++ specifications = [
++ ("ep", 0.003),
++ ("dualprop", 0.025),
++ ("ep", 0.01),
++ ("ff", 0.03),
++ ("ep", 0.03),
++ ("pepita", 0.01),
++ ("fa", 0.003),
++ ("bp", 0.025),
++ ("fa", 0.01),
++ ("dfa", 0.003),
++ ("fa", 0.03),
++ ("dfa", 0.01),
++ ("clean_kp", 0.003),
++ ("dfa", 0.03),
++ ("clean_kp", 0.01),
++ ("sdil", 0.003),
++ ("clean_kp", 0.03),
++ ("sdil", 0.01),
++ ("sdil", 0.03),
++ ]
++ return [p1_command(method, rate) for method, rate in specifications]
++
++
++def p1_command(method, rate):
++ cli_method = {
++ "dualprop": "dualprop-lagr-ff",
++ "clean_kp": "clean-kp",
++ }.get(method, method)
++ runner = "train_ff.py" if method == "ff" else "train.py"
++ name = f"plain-p1-{method}-minicnn-lr{rate_tag(rate)}"
++ command = [
++ sys.executable, runner,
++ "--model", "miniCNN",
++ "--dataset", "cifar10",
++ "--num-epochs", "10",
++ "--batch-size", "100",
++ "--learning-rate", str(rate),
++ "--learning-rate-final", str(rate),
++ "--warmup-learning-rate", str(rate),
++ "--warmup-epochs", "0",
++ "--decay-epochs", "10",
++ "--momentum", "0.9",
++ "--weight-decay", "5e-4",
++ "--dtype", "float32",
++ "--param-dtype", "float32",
++ "--percent-train", "90",
++ "--percent-val", "10",
++ "--seeds", "0",
++ "--feedback-seed", "1729",
++ "--gradient-diagnostics", "none",
++ "--spectral-diagnostics", "none",
++ "--test-policy", "none",
++ "--early-stop-policy", "none",
++ "--learning-algorithm", cli_method,
++ "--experiment-name", name,
++ ]
++ if method == "pepita":
++ command.extend([
++ "--optimizer-schedule", "pepita",
++ "--pepita-projection-scale", "0.05",
++ ])
++ else:
++ command.extend(["--optimizer-schedule", "author"])
++ if method == "ff":
++ command.extend([
++ "--ff-threshold", "2.0",
++ "--ff-score-from-layer", "1",
++ ])
++ if method == "ep":
++ command.extend([
++ "--ep-beta", "0.5",
++ "--ep-dt", "0.5",
++ "--ep-free-steps", "20",
++ "--ep-nudge-steps", "4",
++ ])
++ if method == "dualprop":
++ command.extend([
++ "--loss", "sce",
++ "--alpha", "0.0",
++ "--beta", "0.1",
++ "--inference-sequence", "fwK",
++ "--inference-passes-nudged", "16",
++ ])
++ if method == "sdil":
++ command.extend([
++ "--sdil-traffic-ratio", "4",
++ "--sdil-traffic-seed", "4000",
++ "--sdil-calibration-examples", "64",
++ ])
++ return {
++ "stage": "p1",
++ "method": method,
++ "architecture": "minicnn",
++ "rate": rate,
++ "experiment_name": name,
++ "command": command,
++ }
++
++
++def completed_histogram(experiment_name):
++ paths = glob.glob(os.path.join(
++ ROOT, "runs", experiment_name, "*", "hist.npy"))
++ if not paths:
++ return None
++ if len(paths) != 1:
++ raise RuntimeError(
++ f"expected one result for {experiment_name}, found {len(paths)}")
++ return paths[0]
++
++
++def run_job(job, source, dry_run):
++ result = completed_histogram(job["experiment_name"])
++ if result is not None:
++ print(f"preserving {job['experiment_name']} -> {result}", flush=True)
++ return
++ print("RUN", " ".join(job["command"]), flush=True)
++ if dry_run:
++ return
++ started = time.time()
++ subprocess.run(job["command"], cwd=ROOT, check=True)
++ result = completed_histogram(job["experiment_name"])
++ if result is None:
++ raise RuntimeError(f"job produced no histogram: {job['experiment_name']}")
++ manifest = {
++ **job,
++ "source": source,
++ "histogram": result,
++ "histogram_sha256": sha256(result),
++ "driver_wall_seconds": time.time() - started,
++ "completed_unix_time": time.time(),
++ }
++ manifest_path = os.path.join(
++ os.path.dirname(result), "crossover_manifest.json")
++ with open(manifest_path, "w") as handle:
++ json.dump(manifest, handle, indent=2, sort_keys=True)
++ handle.write("\n")
++ print(f"DONE {job['experiment_name']} -> {manifest_path}", flush=True)
++
++
++def main():
++ parser = argparse.ArgumentParser()
++ parser.add_argument("--stage", choices=("p1",), default="p1")
++ parser.add_argument("--shard-index", type=int, default=0)
++ parser.add_argument("--num-shards", type=int, default=1)
++ parser.add_argument("--method")
++ parser.add_argument("--dry-run", action="store_true")
++ args = parser.parse_args()
++ if not 0 <= args.shard_index < args.num_shards:
++ raise ValueError("invalid shard")
++ source = source_report()
++ jobs = p1_jobs()
++ if args.method:
++ jobs = [job for job in jobs if job["method"] == args.method]
++ jobs = [
++ job for index, job in enumerate(jobs)
++ if index % args.num_shards == args.shard_index
++ ]
++ if not jobs:
++ raise ValueError("no jobs selected")
++ for job in jobs:
++ run_job(job, source, args.dry_run)
++
++
++if __name__ == "__main__":
++ main()
+--
+2.54.0
+
diff --git a/external/dualprop_patches/0012-fix-map-crossover-BP-to-author-CLI.patch b/external/dualprop_patches/0012-fix-map-crossover-BP-to-author-CLI.patch
new file mode 100644
index 0000000..9544fa9
--- /dev/null
+++ b/external/dualprop_patches/0012-fix-map-crossover-BP-to-author-CLI.patch
@@ -0,0 +1,24 @@
+From 8c26ab9f5860003af39fc0a69f5c982906553ab6 Mon Sep 17 00:00:00 2001
+From: YurenHao0426 <Blackhao0426@gmail.com>
+Date: Mon, 27 Jul 2026 13:43:37 -0500
+Subject: [PATCH 12/19] fix: map crossover BP to author CLI
+
+---
+ crossover_grid.py | 1 +
+ 1 file changed, 1 insertion(+)
+
+diff --git a/crossover_grid.py b/crossover_grid.py
+index dae3316..340e174 100644
+--- a/crossover_grid.py
++++ b/crossover_grid.py
+@@ -86,6 +86,7 @@ def p1_jobs():
+
+ def p1_command(method, rate):
+ cli_method = {
++ "bp": "backprop",
+ "dualprop": "dualprop-lagr-ff",
+ "clean_kp": "clean-kp",
+ }.get(method, method)
+--
+2.54.0
+
diff --git a/external/dualprop_patches/0013-analysis-audit-plain-CNN-P1-selector.patch b/external/dualprop_patches/0013-analysis-audit-plain-CNN-P1-selector.patch
new file mode 100644
index 0000000..ef2612d
--- /dev/null
+++ b/external/dualprop_patches/0013-analysis-audit-plain-CNN-P1-selector.patch
@@ -0,0 +1,260 @@
+From 46b74984587f7b76c40f91cfccd3a6939695a121 Mon Sep 17 00:00:00 2001
+From: YurenHao0426 <Blackhao0426@gmail.com>
+Date: Mon, 27 Jul 2026 13:45:04 -0500
+Subject: [PATCH 13/19] analysis: audit plain CNN P1 selector
+
+---
+ analyze_p1.py | 241 ++++++++++++++++++++++++++++++++++++++++++++++++++
+ 1 file changed, 241 insertions(+)
+ create mode 100644 analyze_p1.py
+
+diff --git a/analyze_p1.py b/analyze_p1.py
+new file mode 100644
+index 0000000..05dd52d
+--- /dev/null
++++ b/analyze_p1.py
+@@ -0,0 +1,241 @@
++#!/usr/bin/env python3
++"""Audit and mechanically select the frozen plain-CNN P1 learning rates."""
++import argparse
++import glob
++import hashlib
++import json
++import math
++import os
++
++import numpy as np
++
++
++ROOT = os.path.dirname(os.path.abspath(__file__))
++EXPECTED_SOURCE = {
++ "author_crossover_commit": "90697e8a17ab5e0757a62df40da521766763abe1",
++ "main_commit": "dab65f657ed22ce90fe2b38a360ca62cc8b00e65",
++ "protocol_sha256":
++ "345e0c1245f558a68ea2c212dfa427bdaa30eafd2c8d16c6763691ba354dae61",
++}
++SPECIFICATIONS = (
++ ("bp", 0.025),
++ ("fa", 0.003),
++ ("fa", 0.01),
++ ("fa", 0.03),
++ ("dfa", 0.003),
++ ("dfa", 0.01),
++ ("dfa", 0.03),
++ ("pepita", 0.01),
++ ("ff", 0.03),
++ ("ep", 0.003),
++ ("ep", 0.01),
++ ("ep", 0.03),
++ ("dualprop", 0.025),
++ ("clean_kp", 0.003),
++ ("clean_kp", 0.01),
++ ("clean_kp", 0.03),
++ ("sdil", 0.003),
++ ("sdil", 0.01),
++ ("sdil", 0.03),
++)
++GRIDDED_METHODS = ("fa", "dfa", "ep", "clean_kp", "sdil")
++CLI_METHODS = {
++ "bp": "backprop",
++ "clean_kp": "clean-kp",
++ "dualprop": "dualprop-lagr-ff",
++}
++
++
++def rate_tag(rate):
++ return f"{rate:g}".replace(".", "p")
++
++
++def experiment_name(method, rate):
++ return f"plain-p1-{method}-minicnn-lr{rate_tag(rate)}"
++
++
++def sha256(path):
++ digest = hashlib.sha256()
++ with open(path, "rb") as handle:
++ for chunk in iter(lambda: handle.read(1024 * 1024), b""):
++ digest.update(chunk)
++ return digest.hexdigest()
++
++
++def flag(command, name):
++ try:
++ index = command.index(name)
++ except ValueError as error:
++ raise AssertionError(f"missing command flag {name}") from error
++ if index + 1 >= len(command):
++ raise AssertionError(f"missing value for command flag {name}")
++ return command[index + 1]
++
++
++def scalar(value):
++ return float(np.asarray(value))
++
++
++def audit_histogram(path, method):
++ history = np.load(path, allow_pickle=True).item()
++ if method == "ff":
++ assert history["method"] == "ff"
++ assert history["epochs_per_layer"] == 10
++ assert len(history["layers"]) == history["num_layers"]
++ assert all(
++ len(layer["epochs"]) == 10 for layer in history["layers"])
++ final = history["final"]
++ best_accuracy = float(final["validation_accuracy"])
++ final_accuracy = best_accuracy
++ final_loss = float(history["layers"][-1]["epochs"][-1]["loss"])
++ test_values = (final["test_accuracy"], final["test_time"])
++ epochs_completed = sum(
++ len(layer["epochs"]) for layer in history["layers"])
++ best_epoch = None
++ train_seconds = sum(
++ epoch["runtime"] for layer in history["layers"]
++ for epoch in layer["epochs"])
++ validation_seconds = float(final["validation_time"])
++ else:
++ validation_accuracy = np.asarray(
++ history["val_accuracy"], dtype=float)
++ validation_loss = np.asarray(history["val_loss"], dtype=float)
++ assert validation_accuracy.shape == (10,)
++ assert validation_loss.shape == (10,)
++ assert int(history["epochs_completed"]) == 10
++ assert np.all(np.isfinite(validation_accuracy))
++ best_accuracy = float(np.max(validation_accuracy))
++ final_accuracy = float(validation_accuracy[-1])
++ final_loss = float(validation_loss[-1])
++ assert math.isclose(
++ best_accuracy, scalar(history["best_validation_accuracy"]),
++ rel_tol=0, abs_tol=1e-6)
++ test_values = (
++ history["test_accuracy"], history["test_loss"],
++ history["test_top5accuracy"], history["test_time"])
++ epochs_completed = 10
++ best_epoch = int(history["best_epoch"])
++ train_seconds = float(np.sum(history["train_time"]))
++ validation_seconds = float(np.sum(history["val_time"]))
++ assert all(math.isnan(scalar(value)) for value in test_values)
++ return {
++ "best_validation_accuracy": best_accuracy,
++ "final_validation_accuracy": final_accuracy,
++ "final_validation_loss": final_loss,
++ "epochs_completed": epochs_completed,
++ "best_epoch": best_epoch,
++ "train_seconds": train_seconds,
++ "validation_seconds": validation_seconds,
++ "test_metrics_all_nan": True,
++ }
++
++
++def audit_record(method, rate):
++ name = experiment_name(method, rate)
++ run_root = os.path.join(ROOT, "runs", name)
++ histograms = glob.glob(os.path.join(run_root, "*", "hist.npy"))
++ manifests = glob.glob(
++ os.path.join(run_root, "*", "crossover_manifest.json"))
++ assert len(histograms) == 1, f"{name}: found {len(histograms)} histograms"
++ assert len(manifests) == 1, f"{name}: found {len(manifests)} manifests"
++ histogram = os.path.realpath(histograms[0])
++ with open(manifests[0], encoding="utf-8") as handle:
++ manifest = json.load(handle)
++ assert manifest["stage"] == "p1"
++ assert manifest["method"] == method
++ assert manifest["architecture"] == "minicnn"
++ assert math.isclose(float(manifest["rate"]), rate)
++ assert manifest["experiment_name"] == name
++ assert os.path.realpath(manifest["histogram"]) == histogram
++ assert manifest["histogram_sha256"] == sha256(histogram)
++ for key, expected in EXPECTED_SOURCE.items():
++ assert manifest["source"][key] == expected, f"{name}: source drift"
++ command = manifest["command"]
++ assert flag(command, "--model") == "miniCNN"
++ assert flag(command, "--dataset") == "cifar10"
++ assert flag(command, "--num-epochs") == "10"
++ assert flag(command, "--seeds") == "0"
++ assert flag(command, "--feedback-seed") == "1729"
++ assert flag(command, "--test-policy") == "none"
++ assert flag(command, "--early-stop-policy") == "none"
++ assert flag(command, "--gradient-diagnostics") == "none"
++ assert flag(command, "--spectral-diagnostics") == "none"
++ assert flag(command, "--learning-algorithm") == CLI_METHODS.get(
++ method, method)
++ assert math.isclose(float(flag(command, "--learning-rate")), rate)
++ metrics = audit_histogram(histogram, method)
++ return {
++ "method": method,
++ "rate": rate,
++ "experiment_name": name,
++ "manifest": os.path.relpath(manifests[0], ROOT),
++ "histogram_sha256": manifest["histogram_sha256"],
++ "driver_wall_seconds": manifest.get("driver_wall_seconds"),
++ **metrics,
++ }
++
++
++def selected_candidate(candidates):
++ def rank(candidate):
++ loss = candidate["final_validation_loss"]
++ return (
++ candidate["best_validation_accuracy"],
++ int(math.isfinite(loss)),
++ -candidate["rate"],
++ )
++ return max(candidates, key=rank)
++
++
++def main():
++ parser = argparse.ArgumentParser()
++ parser.add_argument("--output")
++ args = parser.parse_args()
++ expected_names = {
++ experiment_name(method, rate) for method, rate in SPECIFICATIONS}
++ actual_names = {
++ os.path.basename(path)
++ for path in glob.glob(os.path.join(ROOT, "runs", "plain-p1-*"))
++ if glob.glob(os.path.join(path, "*", "hist.npy"))
++ }
++ assert actual_names == expected_names, {
++ "missing": sorted(expected_names - actual_names),
++ "unexpected": sorted(actual_names - expected_names),
++ }
++ records = [
++ audit_record(method, rate) for method, rate in SPECIFICATIONS]
++ selected = {}
++ for method in dict(SPECIFICATIONS):
++ candidates = [
++ record for record in records if record["method"] == method]
++ chosen = (
++ selected_candidate(candidates)
++ if method in GRIDDED_METHODS else candidates[0])
++ selected[method] = {
++ "rate": chosen["rate"],
++ "best_validation_accuracy":
++ chosen["best_validation_accuracy"],
++ "experiment_name": chosen["experiment_name"],
++ }
++ report = {
++ "gate": "pass",
++ "stage": "plain_cnn_p1",
++ "selection_rule":
++ "maximum best validation accuracy, then finite final validation "
++ "loss, then lower learning rate",
++ "expected_source": EXPECTED_SOURCE,
++ "num_expected_records": len(SPECIFICATIONS),
++ "num_audited_records": len(records),
++ "test_policy": "none",
++ "records": records,
++ "selected": selected,
++ }
++ encoded = json.dumps(report, indent=2, sort_keys=True) + "\n"
++ if args.output:
++ with open(args.output, "w", encoding="utf-8") as handle:
++ handle.write(encoded)
++ else:
++ print(encoded, end="")
++
++
++if __name__ == "__main__":
++ main()
+--
+2.54.0
+
diff --git a/external/dualprop_patches/0014-experiment-freeze-plain-CNN-P2-registry.patch b/external/dualprop_patches/0014-experiment-freeze-plain-CNN-P2-registry.patch
new file mode 100644
index 0000000..f6ce895
--- /dev/null
+++ b/external/dualprop_patches/0014-experiment-freeze-plain-CNN-P2-registry.patch
@@ -0,0 +1,228 @@
+From 0645bf8369bf49ca4c23a9b2c19e4742f174addc Mon Sep 17 00:00:00 2001
+From: YurenHao0426 <Blackhao0426@gmail.com>
+Date: Mon, 27 Jul 2026 13:49:15 -0500
+Subject: [PATCH 14/19] experiment: freeze plain CNN P2 registry
+
+---
+ crossover_grid.py | 177 +++++++++++++++++++++++++++++++++++++++++++++-
+ 1 file changed, 175 insertions(+), 2 deletions(-)
+
+diff --git a/crossover_grid.py b/crossover_grid.py
+index 340e174..846a2f1 100644
+--- a/crossover_grid.py
++++ b/crossover_grid.py
+@@ -13,6 +13,14 @@ import time
+ ROOT = os.path.dirname(os.path.abspath(__file__))
+ MAIN_ROOT = "/home/yurenh2/sdil"
+ PROTOCOL = os.path.join(MAIN_ROOT, "PLAIN_CNN_CROSSOVER.md")
++SELECTOR = os.path.join(
++ MAIN_ROOT, "results", "plain_cnn_p1_selector.json")
++SELECTOR_SHA256 = (
++ "3365f4618e0941de7dd7183f69c1f1bdd596a94b60f58de9e4e60ed05c38918c")
++P2_METHODS = (
++ "bp", "fa", "dfa", "pepita", "ff", "ep", "dualprop", "clean_kp",
++ "sdil")
++P2_ARCHITECTURES = ("minicnn", "vgglike", "vgg16")
+
+
+ def output(*args):
+@@ -58,6 +66,18 @@ def rate_tag(rate):
+ return f"{rate:g}".replace(".", "p")
+
+
++def selector_report():
++ if sha256(SELECTOR) != SELECTOR_SHA256:
++ raise RuntimeError("P1 selector artifact hash drift")
++ with open(SELECTOR, encoding="utf-8") as handle:
++ report = json.load(handle)
++ if report["gate"] != "pass" or report["num_audited_records"] != 19:
++ raise RuntimeError("P1 selector gate did not pass")
++ if set(report["selected"]) != set(P2_METHODS):
++ raise RuntimeError("P1 selector method registry drift")
++ return report
++
++
+ def p1_jobs():
+ # Expensive jobs are deliberately interleaved across modulo shards.
+ specifications = [
+@@ -84,6 +104,159 @@ def p1_jobs():
+ return [p1_command(method, rate) for method, rate in specifications]
+
+
++def p2_jobs():
++ report = selector_report()
++ # Expensive largest-model cells start first. Modulo-2 assignment places
++ # VGG16 EP and DP on different physical GPUs.
++ schedule = [
++ ("ep", "vgg16"),
++ ("dualprop", "vgg16"),
++ ("ff", "vgg16"),
++ ("pepita", "vgg16"),
++ ("sdil", "vgg16"),
++ ("clean_kp", "vgg16"),
++ ("fa", "vgg16"),
++ ("dfa", "vgg16"),
++ ("bp", "vgg16"),
++ ("ep", "vgglike"),
++ ("dualprop", "vgglike"),
++ ("ff", "vgglike"),
++ ("pepita", "vgglike"),
++ ("sdil", "vgglike"),
++ ("clean_kp", "vgglike"),
++ ("fa", "vgglike"),
++ ("dfa", "vgglike"),
++ ("bp", "vgglike"),
++ ("ep", "minicnn"),
++ ("dualprop", "minicnn"),
++ ("ff", "minicnn"),
++ ("pepita", "minicnn"),
++ ("sdil", "minicnn"),
++ ("clean_kp", "minicnn"),
++ ("fa", "minicnn"),
++ ("dfa", "minicnn"),
++ ("bp", "minicnn"),
++ ]
++ assert len(schedule) == len(P2_METHODS) * len(P2_ARCHITECTURES)
++ assert set(schedule) == {
++ (method, architecture)
++ for method in P2_METHODS for architecture in P2_ARCHITECTURES}
++ return [
++ p2_command(
++ method, architecture, report["selected"][method]["rate"])
++ for method, architecture in schedule
++ ]
++
++
++def p2_command(method, architecture, rate):
++ cli_method = {
++ "bp": "backprop",
++ "dualprop": "dualprop-lagr-ff",
++ "clean_kp": "clean-kp",
++ }.get(method, method)
++ model = {
++ "minicnn": "miniCNN",
++ "vgglike": "VGGlike",
++ "vgg16": "VGG16",
++ }[architecture]
++ runner = "train_ff.py" if method == "ff" else "train.py"
++ if method in ("bp", "fa", "dfa", "dualprop", "clean_kp", "sdil"):
++ epochs = 130
++ final_rate = 2e-6
++ warmup_rate = 0.001
++ warmup_epochs = 10
++ decay_epochs = 120
++ optimizer_schedule = "author"
++ elif method == "pepita":
++ epochs = 100
++ final_rate = warmup_rate = rate
++ warmup_epochs = 0
++ decay_epochs = epochs
++ optimizer_schedule = "pepita"
++ elif method == "ff":
++ epochs = 40
++ final_rate = warmup_rate = rate
++ warmup_epochs = 0
++ decay_epochs = epochs
++ optimizer_schedule = "author"
++ elif method == "ep":
++ epochs = 100
++ final_rate = warmup_rate = rate
++ warmup_epochs = 0
++ decay_epochs = epochs
++ optimizer_schedule = "author"
++ else:
++ raise ValueError(method)
++ name = f"plain-p2-{method}-{architecture}"
++ command = [
++ sys.executable, runner,
++ "--model", model,
++ "--dataset", "cifar10",
++ "--num-epochs", str(epochs),
++ "--batch-size", "100",
++ "--learning-rate", str(rate),
++ "--learning-rate-final", str(final_rate),
++ "--warmup-learning-rate", str(warmup_rate),
++ "--warmup-epochs", str(warmup_epochs),
++ "--decay-epochs", str(decay_epochs),
++ "--momentum", "0.9",
++ "--weight-decay", "5e-4",
++ "--dtype", "float32",
++ "--param-dtype", "float32",
++ "--percent-train", "90",
++ "--percent-val", "10",
++ "--seeds", "0",
++ "--feedback-seed", "1729",
++ "--gradient-diagnostics", "none",
++ "--spectral-diagnostics", "none",
++ "--test-policy", "none",
++ "--early-stop-policy", "none",
++ "--learning-algorithm", cli_method,
++ "--experiment-name", name,
++ "--optimizer-schedule", optimizer_schedule,
++ ]
++ if method == "pepita":
++ command.extend(["--pepita-projection-scale", "0.05"])
++ if method == "ff":
++ command.extend([
++ "--ff-threshold", "2.0",
++ "--ff-score-from-layer", "1",
++ ])
++ if method == "ep":
++ command.extend([
++ "--ep-beta", "0.5",
++ "--ep-dt", "0.5",
++ "--ep-free-steps", "20",
++ "--ep-nudge-steps", "4",
++ ])
++ if method == "dualprop":
++ command.extend([
++ "--loss", "sce",
++ "--alpha", "0.0",
++ "--beta", "0.1",
++ "--inference-sequence", "fwK",
++ "--inference-passes-nudged", "16",
++ ])
++ if method == "sdil":
++ command.extend([
++ "--sdil-traffic-ratio", "4",
++ "--sdil-traffic-seed", "4000",
++ "--sdil-calibration-examples", "64",
++ ])
++ return {
++ "stage": "p2",
++ "method": method,
++ "architecture": architecture,
++ "rate": rate,
++ "expected_epochs": epochs,
++ "experiment_name": name,
++ "selector_path": SELECTOR,
++ "selector_sha256": SELECTOR_SHA256,
++ "timeout_seconds": 48 * 60 * 60,
++ "command": command,
++ }
++
++
+ def p1_command(method, rate):
+ cli_method = {
+ "bp": "backprop",
+@@ -203,7 +376,7 @@ def run_job(job, source, dry_run):
+
+ def main():
+ parser = argparse.ArgumentParser()
+- parser.add_argument("--stage", choices=("p1",), default="p1")
++ parser.add_argument("--stage", choices=("p1", "p2"), default="p1")
+ parser.add_argument("--shard-index", type=int, default=0)
+ parser.add_argument("--num-shards", type=int, default=1)
+ parser.add_argument("--method")
+@@ -212,7 +385,7 @@ def main():
+ if not 0 <= args.shard_index < args.num_shards:
+ raise ValueError("invalid shard")
+ source = source_report()
+- jobs = p1_jobs()
++ jobs = p1_jobs() if args.stage == "p1" else p2_jobs()
+ if args.method:
+ jobs = [job for job in jobs if job["method"] == args.method]
+ jobs = [
+--
+2.54.0
+
diff --git a/external/dualprop_patches/0015-experiment-audit-P2-execution-and-resource-use.patch b/external/dualprop_patches/0015-experiment-audit-P2-execution-and-resource-use.patch
new file mode 100644
index 0000000..8a019cd
--- /dev/null
+++ b/external/dualprop_patches/0015-experiment-audit-P2-execution-and-resource-use.patch
@@ -0,0 +1,341 @@
+From 0e868b670e05e123a523ea18986b3e1423e11726 Mon Sep 17 00:00:00 2001
+From: YurenHao0426 <Blackhao0426@gmail.com>
+Date: Mon, 27 Jul 2026 13:50:34 -0500
+Subject: [PATCH 15/19] experiment: audit P2 execution and resource use
+
+---
+ crossover_grid.py | 289 +++++++++++++++++++++++++++++++++++++++++++++-
+ 1 file changed, 288 insertions(+), 1 deletion(-)
+
+diff --git a/crossover_grid.py b/crossover_grid.py
+index 846a2f1..cbd9b98 100644
+--- a/crossover_grid.py
++++ b/crossover_grid.py
+@@ -5,6 +5,7 @@ import glob
+ import hashlib
+ import json
+ import os
++import signal
+ import subprocess
+ import sys
+ import time
+@@ -37,6 +38,30 @@ def sha256(path):
+ return digest.hexdigest()
+
+
++def sha256_text(value):
++ return hashlib.sha256(value.encode("utf-8")).hexdigest()
++
++
++def hash_tree(path):
++ digest = hashlib.sha256()
++ files = []
++ for directory, _, names in os.walk(path):
++ for name in sorted(names):
++ files.append(os.path.join(directory, name))
++ for filename in sorted(files):
++ relative = os.path.relpath(filename, path)
++ digest.update(relative.encode("utf-8") + b"\0")
++ with open(filename, "rb") as handle:
++ for chunk in iter(lambda: handle.read(1024 * 1024), b""):
++ digest.update(chunk)
++ return {
++ "path": os.path.realpath(path),
++ "sha256": digest.hexdigest(),
++ "num_files": len(files),
++ "bytes": sum(os.path.getsize(filename) for filename in files),
++ }
++
++
+ def source_report():
+ commit = output("git", "rev-parse", "HEAD")
+ dirty = output(
+@@ -52,14 +77,21 @@ def source_report():
+ text=True).stdout.strip()
+ if main_dirty:
+ raise RuntimeError("crossover runs require clean tracked main source")
+- return {
++ pip_freeze = output(sys.executable, "-m", "pip", "freeze")
++ report = {
+ "author_crossover_commit": commit,
+ "main_commit": main_commit,
+ "protocol_path": PROTOCOL,
+ "protocol_sha256": sha256(PROTOCOL),
+ "python": sys.executable,
+ "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"),
++ "pip_freeze_sha256": sha256_text(pip_freeze),
++ "python_version": sys.version,
+ }
++ dataset = os.path.join(ROOT, "data", "cifar10", "3.0.2")
++ if os.path.isdir(dataset):
++ report["cifar10_tfds_tree"] = hash_tree(dataset)
++ return report
+
+
+ def rate_tag(rate):
+@@ -345,7 +377,262 @@ def completed_histogram(experiment_name):
+ return paths[0]
+
+
++def physical_gpu_report():
++ index = os.environ.get("CUDA_VISIBLE_DEVICES")
++ if index not in ("5", "7"):
++ raise RuntimeError(
++ "formal P2 is authorized only on physical GPUs 5 and 7")
++ query = output(
++ "nvidia-smi",
++ "--query-gpu=index,uuid,name,memory.total",
++ "--format=csv,noheader,nounits")
++ rows = [row.split(", ", 3) for row in query.splitlines()]
++ matches = [row for row in rows if row[0] == index]
++ if len(matches) != 1:
++ raise RuntimeError(f"cannot resolve physical GPU {index}")
++ physical_index, uuid, name, memory_mib = matches[0]
++ return {
++ "physical_index": int(physical_index),
++ "uuid": uuid,
++ "name": name,
++ "total_memory_mib": int(memory_mib),
++ }
++
++
++def descendant_pids(root_pid):
++ listing = subprocess.run(
++ ["ps", "-eo", "pid=,ppid="], check=True, capture_output=True,
++ text=True).stdout
++ children = {}
++ for row in listing.splitlines():
++ pid, parent = (int(value) for value in row.split())
++ children.setdefault(parent, []).append(pid)
++ descendants = {root_pid}
++ frontier = [root_pid]
++ while frontier:
++ parent = frontier.pop()
++ for child in children.get(parent, ()):
++ if child not in descendants:
++ descendants.add(child)
++ frontier.append(child)
++ return descendants
++
++
++def process_gpu_memory_mib(root_pid, gpu_uuid):
++ query = subprocess.run(
++ [
++ "nvidia-smi",
++ "--query-compute-apps=pid,gpu_uuid,used_memory",
++ "--format=csv,noheader,nounits",
++ ],
++ check=True, capture_output=True, text=True).stdout
++ descendants = descendant_pids(root_pid)
++ total = 0
++ for row in query.splitlines():
++ pid, uuid, memory = row.split(", ")
++ if uuid == gpu_uuid and int(pid) in descendants:
++ total += int(memory)
++ return total
++
++
++def architecture_work(architecture):
++ specifications = {
++ "minicnn": (
++ [64, 64], [True, True], [10]),
++ "vgglike": (
++ [128, 256, 512, 512], [True, True, True, True], [10]),
++ "vgg16": (
++ [64, 64, 128, 128, 256, 256, 256, 512, 512, 512, 512,
++ 512, 512],
++ [False, True, False, True, False, False, True, False,
++ False, True, False, False, True],
++ [4096, 4096, 10]),
++ }
++ features, pooling, dense = specifications[architecture]
++ height = width = 32
++ channels = 3
++ parameters = 0
++ forward_macs = 0
++ for output_channels, pool in zip(features, pooling):
++ forward_macs += (
++ height * width * channels * output_channels * 3 * 3)
++ parameters += 3 * 3 * channels * output_channels + output_channels
++ channels = output_channels
++ if pool:
++ height //= 2
++ width //= 2
++ inputs = height * width * channels
++ for outputs in dense:
++ forward_macs += inputs * outputs
++ parameters += inputs * outputs + outputs
++ inputs = outputs
++ return {
++ "forward_parameter_count": parameters,
++ "forward_affine_macs_per_example": forward_macs,
++ "num_trainable_layers": len(features) + len(dense),
++ }
++
++
++def p2_work_report(job):
++ epochs = job["expected_epochs"]
++ method = job["method"]
++ architecture = architecture_work(job["architecture"])
++ train_examples = 45000 * epochs
++ validation_examples = 5000 if method == "ff" else 5000 * epochs
++ layers = architecture["num_trainable_layers"]
++ presentations = train_examples
++ if method == "pepita":
++ presentations *= 2
++ elif method == "ff":
++ presentations *= 2 * layers
++ feedforward_train_multipliers = {
++ "bp": 2,
++ "fa": 2,
++ "dfa": 2,
++ "pepita": 3,
++ "ff": 2 * layers,
++ "ep": 1,
++ "dualprop": 1,
++ "clean_kp": 2,
++ "sdil": 2,
++ }
++ relaxation_passes = 0
++ if method == "ep":
++ relaxation_passes = train_examples * (20 + 4)
++ relaxation_passes += validation_examples * 20
++ elif method == "dualprop":
++ relaxation_passes = train_examples * 16
++ candidate_evaluations = 10 * validation_examples if method == "ff" else 0
++ feedforward_passes = (
++ train_examples * feedforward_train_multipliers[method]
++ + (0 if method == "ff" else validation_examples)
++ + candidate_evaluations)
++ return {
++ **architecture,
++ "ordinary_training_examples": train_examples,
++ "ordinary_validation_examples": validation_examples,
++ "training_example_presentations": presentations,
++ "feedforward_example_passes": feedforward_passes,
++ "relaxation_example_passes": relaxation_passes,
++ "candidate_label_evaluation_presentations": candidate_evaluations,
++ "task_loss_query_batches": (
++ 0 if method == "ff" else 450 * epochs),
++ "neutral_predictor_observations": 64 if method == "sdil" else 0,
++ "counting_note": (
++ "Explicit example/pass counts; affine correlation work is "
++ "audited separately and is not inferred from wall time."),
++ }
++
++
++def record_path(experiment_name):
++ paths = glob.glob(os.path.join(
++ ROOT, "runs", experiment_name, "*", "crossover_manifest.json"))
++ if len(paths) > 1:
++ raise RuntimeError(
++ f"expected at most one record for {experiment_name}")
++ return paths[0] if paths else None
++
++
++def result_directory(experiment_name, create_if_missing=False):
++ root = os.path.join(ROOT, "runs", experiment_name)
++ paths = [
++ path for path in glob.glob(
++ os.path.join(root, "*"))
++ if os.path.isdir(path)
++ ]
++ if not paths and create_if_missing:
++ directory = os.path.join(
++ root, "driver_" + time.strftime("%Y_%m_%d_%H_%M_%S"))
++ os.makedirs(directory)
++ paths = [directory]
++ if len(paths) != 1:
++ raise RuntimeError(
++ f"expected one run directory for {experiment_name}, "
++ f"found {len(paths)}")
++ return paths[0]
++
++
++def terminate_process_group(process):
++ try:
++ os.killpg(process.pid, signal.SIGTERM)
++ process.wait(timeout=10)
++ except subprocess.TimeoutExpired:
++ os.killpg(process.pid, signal.SIGKILL)
++ process.wait()
++
++
++def run_p2_job(job, source, dry_run):
++ existing = record_path(job["experiment_name"])
++ if existing is not None:
++ print(f"preserving {job['experiment_name']} -> {existing}", flush=True)
++ return
++ run_root = os.path.join(ROOT, "runs", job["experiment_name"])
++ if os.path.isdir(run_root):
++ raise RuntimeError(
++ f"orphaned P2 directory requires audit: {run_root}")
++ print("RUN", " ".join(job["command"]), flush=True)
++ if dry_run:
++ return
++ gpu = physical_gpu_report()
++ started = time.time()
++ process = subprocess.Popen(
++ job["command"], cwd=ROOT, start_new_session=True)
++ peak_memory_mib = 0
++ timed_out = False
++ while process.poll() is None:
++ elapsed = time.time() - started
++ if elapsed > job["timeout_seconds"]:
++ timed_out = True
++ terminate_process_group(process)
++ break
++ try:
++ peak_memory_mib = max(
++ peak_memory_mib,
++ process_gpu_memory_mib(process.pid, gpu["uuid"]))
++ except (OSError, subprocess.SubprocessError, ValueError):
++ pass
++ time.sleep(1)
++ return_code = process.returncode
++ elapsed = time.time() - started
++ histogram = completed_histogram(job["experiment_name"])
++ if timed_out:
++ status = "timeout"
++ elif return_code != 0:
++ status = "nonzero_exit"
++ elif histogram is None:
++ status = "missing_histogram"
++ else:
++ status = "completed"
++ directory = result_directory(
++ job["experiment_name"], create_if_missing=True)
++ manifest = {
++ **job,
++ "source": source,
++ "status": status,
++ "return_code": return_code,
++ "histogram": histogram,
++ "histogram_sha256": sha256(histogram) if histogram else None,
++ "driver_wall_seconds": elapsed,
++ "completed_unix_time": time.time(),
++ "hardware": {
++ **gpu,
++ "peak_process_gpu_memory_mib": peak_memory_mib,
++ "memory_sampling_period_seconds": 1,
++ },
++ "work": p2_work_report(job),
++ }
++ manifest_path = os.path.join(directory, "crossover_manifest.json")
++ with open(manifest_path, "w") as handle:
++ json.dump(manifest, handle, indent=2, sort_keys=True)
++ handle.write("\n")
++ print(
++ f"DONE status={status} {job['experiment_name']} -> {manifest_path}",
++ flush=True)
++
++
+ def run_job(job, source, dry_run):
++ if job["stage"] == "p2":
++ return run_p2_job(job, source, dry_run)
+ result = completed_histogram(job["experiment_name"])
+ if result is not None:
+ print(f"preserving {job['experiment_name']} -> {result}", flush=True)
+--
+2.54.0
+
diff --git a/external/dualprop_patches/0016-experiment-lock-P2-source-and-cell-registry.patch b/external/dualprop_patches/0016-experiment-lock-P2-source-and-cell-registry.patch
new file mode 100644
index 0000000..47fd5b6
--- /dev/null
+++ b/external/dualprop_patches/0016-experiment-lock-P2-source-and-cell-registry.patch
@@ -0,0 +1,70 @@
+From 18b17f442b045686bee420b81b3e889cb80fd71e Mon Sep 17 00:00:00 2001
+From: YurenHao0426 <Blackhao0426@gmail.com>
+Date: Mon, 27 Jul 2026 13:51:07 -0500
+Subject: [PATCH 16/19] experiment: lock P2 source and cell registry
+
+---
+ crossover_grid.py | 40 ++++++++++++++++++++++++++++++++++++++++
+ 1 file changed, 40 insertions(+)
+
+diff --git a/crossover_grid.py b/crossover_grid.py
+index cbd9b98..4523250 100644
+--- a/crossover_grid.py
++++ b/crossover_grid.py
+@@ -180,6 +180,43 @@ def p2_jobs():
+ ]
+
+
++def p2_registry_sha256(jobs):
++ encoded = json.dumps(
++ jobs, sort_keys=True, separators=(",", ":")).encode("utf-8")
++ return hashlib.sha256(encoded).hexdigest()
++
++
++def ensure_p2_launch(source, jobs):
++ path = os.path.join(ROOT, "runs", "plain-p2-launch.json")
++ immutable = {
++ "stage": "p2",
++ "source": source,
++ "selector_path": SELECTOR,
++ "selector_sha256": SELECTOR_SHA256,
++ "num_registered_cells": len(jobs),
++ "registry_sha256": p2_registry_sha256(jobs),
++ "jobs": jobs,
++ }
++ if os.path.exists(path):
++ with open(path, encoding="utf-8") as handle:
++ existing = json.load(handle)
++ for key, value in immutable.items():
++ if existing.get(key) != value:
++ raise RuntimeError(f"P2 launch lock drift in {key}")
++ return path
++ record = {**immutable, "created_unix_time": time.time()}
++ os.makedirs(os.path.dirname(path), exist_ok=True)
++ try:
++ descriptor = os.open(
++ path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o644)
++ except FileExistsError:
++ return ensure_p2_launch(source, jobs)
++ with os.fdopen(descriptor, "w", encoding="utf-8") as handle:
++ json.dump(record, handle, indent=2, sort_keys=True)
++ handle.write("\n")
++ return path
++
++
+ def p2_command(method, architecture, rate):
+ cli_method = {
+ "bp": "backprop",
+@@ -673,6 +710,9 @@ def main():
+ raise ValueError("invalid shard")
+ source = source_report()
+ jobs = p1_jobs() if args.stage == "p1" else p2_jobs()
++ if args.stage == "p2" and not args.dry_run:
++ launch = ensure_p2_launch(source, jobs)
++ print(f"P2 launch lock: {launch}", flush=True)
+ if args.method:
+ jobs = [job for job in jobs if job["method"] == args.method]
+ jobs = [
+--
+2.54.0
+
diff --git a/external/dualprop_patches/0017-analysis-freeze-complete-P2-audit-gate.patch b/external/dualprop_patches/0017-analysis-freeze-complete-P2-audit-gate.patch
new file mode 100644
index 0000000..2119d99
--- /dev/null
+++ b/external/dualprop_patches/0017-analysis-freeze-complete-P2-audit-gate.patch
@@ -0,0 +1,248 @@
+From 52f1b56e3a3fc80cc25b9ac9efc935f8d135961e Mon Sep 17 00:00:00 2001
+From: YurenHao0426 <Blackhao0426@gmail.com>
+Date: Mon, 27 Jul 2026 13:51:49 -0500
+Subject: [PATCH 17/19] analysis: freeze complete P2 audit gate
+
+---
+ analyze_p2.py | 229 ++++++++++++++++++++++++++++++++++++++++++++++++++
+ 1 file changed, 229 insertions(+)
+ create mode 100644 analyze_p2.py
+
+diff --git a/analyze_p2.py b/analyze_p2.py
+new file mode 100644
+index 0000000..d997ec5
+--- /dev/null
++++ b/analyze_p2.py
+@@ -0,0 +1,229 @@
++#!/usr/bin/env python3
++"""Audit the complete frozen 27-cell plain-CNN P2 validation panel."""
++import argparse
++import glob
++import hashlib
++import json
++import math
++import os
++
++import numpy as np
++
++import crossover_grid
++
++
++ROOT = os.path.dirname(os.path.abspath(__file__))
++STATUSES = ("completed", "nonzero_exit", "timeout", "missing_histogram")
++
++
++def sha256(path):
++ digest = hashlib.sha256()
++ with open(path, "rb") as handle:
++ for chunk in iter(lambda: handle.read(1024 * 1024), b""):
++ digest.update(chunk)
++ return digest.hexdigest()
++
++
++def scalar(value):
++ return float(np.asarray(value))
++
++
++def all_nan(values):
++ return all(math.isnan(scalar(value)) for value in values)
++
++
++def history_metrics(path, job):
++ history = np.load(path, allow_pickle=True).item()
++ method = job["method"]
++ expected_epochs = job["expected_epochs"]
++ if method == "ff":
++ assert history["method"] == "ff"
++ assert history["epochs_per_layer"] == expected_epochs
++ assert len(history["layers"]) == history["num_layers"]
++ layer_epochs = [
++ len(layer["epochs"]) for layer in history["layers"]]
++ assert all(count == expected_epochs for count in layer_epochs)
++ final = history["final"]
++ assert all_nan((final["test_accuracy"], final["test_time"]))
++ layer_losses = [
++ float(epoch["loss"]) for layer in history["layers"]
++ for epoch in layer["epochs"]]
++ nonfinite = [
++ index + 1 for index, value in enumerate(layer_losses)
++ if not math.isfinite(value)]
++ validation_accuracy = float(final["validation_accuracy"])
++ return {
++ "outcome": "nonfinite" if nonfinite else "finite",
++ "best_validation_accuracy": validation_accuracy,
++ "final_validation_accuracy": validation_accuracy,
++ "final_validation_loss": None,
++ "best_epoch": None,
++ "epochs_completed": sum(layer_epochs),
++ "first_nonfinite_step": nonfinite[0] if nonfinite else None,
++ "train_seconds": sum(
++ float(epoch["runtime"]) for layer in history["layers"]
++ for epoch in layer["epochs"]),
++ "validation_seconds": float(final["validation_time"]),
++ "test_metrics_all_nan": True,
++ }
++ completed = int(history["epochs_completed"])
++ assert 1 <= completed <= expected_epochs
++ validation_accuracy = np.asarray(
++ history["val_accuracy"], dtype=float)[:completed]
++ validation_loss = np.asarray(
++ history["val_loss"], dtype=float)[:completed]
++ train_loss = np.asarray(history["train_loss"], dtype=float)[:completed]
++ assert np.asarray(history["val_accuracy"]).shape == (expected_epochs,)
++ assert np.asarray(history["train_loss"]).shape == (expected_epochs,)
++ assert all_nan((
++ history["test_accuracy"], history["test_loss"],
++ history["test_top5accuracy"], history["test_time"]))
++ finite_accuracy = validation_accuracy[np.isfinite(validation_accuracy)]
++ best_accuracy = (
++ float(np.max(finite_accuracy)) if finite_accuracy.size else None)
++ nonfinite_mask = (
++ ~np.isfinite(validation_accuracy)
++ | ~np.isfinite(validation_loss)
++ | ~np.isfinite(train_loss))
++ nonfinite_indices = np.flatnonzero(nonfinite_mask)
++ first_nonfinite = (
++ int(nonfinite_indices[0]) + 1 if nonfinite_indices.size else None)
++ if first_nonfinite is None and completed != expected_epochs:
++ outcome = "premature_finite_stop"
++ elif first_nonfinite is None:
++ outcome = "finite"
++ else:
++ outcome = "nonfinite"
++ if best_accuracy is not None:
++ assert math.isclose(
++ best_accuracy, scalar(history["best_validation_accuracy"]),
++ rel_tol=0, abs_tol=1e-5)
++ return {
++ "outcome": outcome,
++ "best_validation_accuracy": best_accuracy,
++ "final_validation_accuracy": float(validation_accuracy[-1]),
++ "final_validation_loss": float(validation_loss[-1]),
++ "best_epoch": int(history["best_epoch"]),
++ "epochs_completed": completed,
++ "first_nonfinite_step": first_nonfinite,
++ "train_seconds": float(
++ np.sum(np.asarray(history["train_time"])[:completed])),
++ "validation_seconds": float(
++ np.sum(np.asarray(history["val_time"])[:completed])),
++ "test_metrics_all_nan": True,
++ }
++
++
++def record_manifest(experiment_name):
++ paths = glob.glob(os.path.join(
++ ROOT, "runs", experiment_name, "*", "crossover_manifest.json"))
++ assert len(paths) == 1, (
++ f"{experiment_name}: expected one manifest, found {len(paths)}")
++ return paths[0]
++
++
++def audit_record(job, launch):
++ manifest_path = record_manifest(job["experiment_name"])
++ with open(manifest_path, encoding="utf-8") as handle:
++ manifest = json.load(handle)
++ for key in (
++ "stage", "method", "architecture", "rate", "expected_epochs",
++ "experiment_name", "selector_path", "selector_sha256",
++ "timeout_seconds", "command"):
++ assert manifest[key] == job[key], (
++ f"{job['experiment_name']}: registry drift in {key}")
++ assert manifest["source"] == launch["source"]
++ assert manifest["status"] in STATUSES
++ assert manifest["selector_sha256"] == crossover_grid.SELECTOR_SHA256
++ assert manifest["work"] == crossover_grid.p2_work_report(job)
++ hardware = manifest["hardware"]
++ assert hardware["physical_index"] in (5, 7)
++ assert hardware["name"] == "NVIDIA GeForce GTX 1080"
++ assert hardware["peak_process_gpu_memory_mib"] >= 0
++ histogram = manifest["histogram"]
++ if histogram is None:
++ assert manifest["histogram_sha256"] is None
++ metrics = {
++ "outcome": manifest["status"],
++ "best_validation_accuracy": None,
++ "final_validation_accuracy": None,
++ "final_validation_loss": None,
++ "best_epoch": None,
++ "epochs_completed": 0,
++ "first_nonfinite_step": None,
++ "train_seconds": None,
++ "validation_seconds": None,
++ "test_metrics_all_nan": None,
++ }
++ else:
++ assert os.path.isfile(histogram)
++ assert manifest["histogram_sha256"] == sha256(histogram)
++ metrics = history_metrics(histogram, job)
++ return {
++ "method": job["method"],
++ "architecture": job["architecture"],
++ "rate": job["rate"],
++ "status": manifest["status"],
++ "manifest": os.path.relpath(manifest_path, ROOT),
++ "histogram_sha256": manifest["histogram_sha256"],
++ "driver_wall_seconds": manifest["driver_wall_seconds"],
++ "physical_gpu": hardware,
++ "work": manifest["work"],
++ **metrics,
++ }
++
++
++def main():
++ parser = argparse.ArgumentParser()
++ parser.add_argument("--output")
++ parser.add_argument("--allow-partial", action="store_true")
++ args = parser.parse_args()
++ launch_path = os.path.join(ROOT, "runs", "plain-p2-launch.json")
++ with open(launch_path, encoding="utf-8") as handle:
++ launch = json.load(handle)
++ jobs = crossover_grid.p2_jobs()
++ assert launch["stage"] == "p2"
++ assert launch["num_registered_cells"] == 27
++ assert launch["registry_sha256"] == crossover_grid.p2_registry_sha256(
++ jobs)
++ assert launch["jobs"] == jobs
++ assert launch["selector_sha256"] == crossover_grid.SELECTOR_SHA256
++ assert launch["source"]["cifar10_tfds_tree"]["num_files"] == 5
++ expected = {job["experiment_name"]: job for job in jobs}
++ actual = {
++ os.path.basename(path)
++ for path in glob.glob(os.path.join(ROOT, "runs", "plain-p2-*"))
++ if os.path.isdir(path) and glob.glob(
++ os.path.join(path, "*", "crossover_manifest.json"))
++ }
++ unexpected = actual - set(expected)
++ missing = set(expected) - actual
++ assert not unexpected, f"unexpected P2 cells: {sorted(unexpected)}"
++ if missing and not args.allow_partial:
++ raise AssertionError(f"missing P2 cells: {sorted(missing)}")
++ records = [
++ audit_record(job, launch) for job in jobs
++ if job["experiment_name"] in actual]
++ report = {
++ "gate": "pass" if not missing else "partial",
++ "stage": "plain_cnn_p2",
++ "launch_record": os.path.relpath(launch_path, ROOT),
++ "source": launch["source"],
++ "selector_sha256": launch["selector_sha256"],
++ "registry_sha256": launch["registry_sha256"],
++ "num_expected_records": len(jobs),
++ "num_audited_records": len(records),
++ "missing_experiments": sorted(missing),
++ "test_policy": "none",
++ "records": records,
++ }
++ encoded = json.dumps(report, indent=2, sort_keys=True) + "\n"
++ if args.output:
++ with open(args.output, "w", encoding="utf-8") as handle:
++ handle.write(encoded)
++ else:
++ print(encoded, end="")
++
++
++if __name__ == "__main__":
++ main()
+--
+2.54.0
+
diff --git a/external/dualprop_patches/0018-fix-separate-P2-source-from-cell-hardware.patch b/external/dualprop_patches/0018-fix-separate-P2-source-from-cell-hardware.patch
new file mode 100644
index 0000000..548bf3a
--- /dev/null
+++ b/external/dualprop_patches/0018-fix-separate-P2-source-from-cell-hardware.patch
@@ -0,0 +1,24 @@
+From de509a8d50f657d7df43b721e5366f3dfd9a2d6d Mon Sep 17 00:00:00 2001
+From: YurenHao0426 <Blackhao0426@gmail.com>
+Date: Mon, 27 Jul 2026 13:52:39 -0500
+Subject: [PATCH 18/19] fix: separate P2 source from cell hardware
+
+---
+ crossover_grid.py | 1 -
+ 1 file changed, 1 deletion(-)
+
+diff --git a/crossover_grid.py b/crossover_grid.py
+index 4523250..bb6d9d5 100644
+--- a/crossover_grid.py
++++ b/crossover_grid.py
+@@ -84,7 +84,6 @@ def source_report():
+ "protocol_path": PROTOCOL,
+ "protocol_sha256": sha256(PROTOCOL),
+ "python": sys.executable,
+- "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"),
+ "pip_freeze_sha256": sha256_text(pip_freeze),
+ "python_version": sys.version,
+ }
+--
+2.54.0
+
diff --git a/external/dualprop_patches/0019-experiment-add-portable-A6000-P2-profile.patch b/external/dualprop_patches/0019-experiment-add-portable-A6000-P2-profile.patch
new file mode 100644
index 0000000..d5e0542
--- /dev/null
+++ b/external/dualprop_patches/0019-experiment-add-portable-A6000-P2-profile.patch
@@ -0,0 +1,217 @@
+From c88ebbc422cb7ff8ae117d924e6fa28c41729022 Mon Sep 17 00:00:00 2001
+From: YurenHao0426 <Blackhao0426@gmail.com>
+Date: Thu, 30 Jul 2026 16:31:36 -0500
+Subject: [PATCH 19/19] experiment: add portable A6000 P2 profile
+
+---
+ analyze_p2.py | 16 +++++++++--
+ crossover_grid.py | 71 +++++++++++++++++++++++++++++++++++++++--------
+ 2 files changed, 74 insertions(+), 13 deletions(-)
+
+diff --git a/analyze_p2.py b/analyze_p2.py
+index d997ec5..696fef0 100644
+--- a/analyze_p2.py
++++ b/analyze_p2.py
+@@ -137,8 +137,18 @@ def audit_record(job, launch):
+ assert manifest["selector_sha256"] == crossover_grid.SELECTOR_SHA256
+ assert manifest["work"] == crossover_grid.p2_work_report(job)
+ hardware = manifest["hardware"]
+- assert hardware["physical_index"] in (5, 7)
+- assert hardware["name"] == "NVIDIA GeForce GTX 1080"
++ hardware_policy = launch["hardware_policy"]
++ assert hardware["hardware_profile"] == hardware_policy["profile"]
++ assert (
++ hardware["physical_index"]
++ in hardware_policy["allowed_physical_gpu_indices"]
++ )
++ assert (
++ hardware["cuda_visible_devices"]
++ == str(hardware["physical_index"])
++ )
++ assert hardware["uuid"]
++ assert hardware["name"] in hardware_policy["allowed_device_names"]
+ assert hardware["peak_process_gpu_memory_mib"] >= 0
+ histogram = manifest["histogram"]
+ if histogram is None:
+@@ -188,6 +198,8 @@ def main():
+ jobs)
+ assert launch["jobs"] == jobs
+ assert launch["selector_sha256"] == crossover_grid.SELECTOR_SHA256
++ assert launch["hardware_policy"]["profile"] in (
++ "timan107_gtx1080", "a6000")
+ assert launch["source"]["cifar10_tfds_tree"]["num_files"] == 5
+ expected = {job["experiment_name"]: job for job in jobs}
+ actual = {
+diff --git a/crossover_grid.py b/crossover_grid.py
+index bb6d9d5..1f14d21 100644
+--- a/crossover_grid.py
++++ b/crossover_grid.py
+@@ -12,7 +12,7 @@ import time
+
+
+ ROOT = os.path.dirname(os.path.abspath(__file__))
+-MAIN_ROOT = "/home/yurenh2/sdil"
++MAIN_ROOT = os.environ.get("SDIL_MAIN_ROOT", "/home/yurenh2/sdil")
+ PROTOCOL = os.path.join(MAIN_ROOT, "PLAIN_CNN_CROSSOVER.md")
+ SELECTOR = os.path.join(
+ MAIN_ROOT, "results", "plain_cnn_p1_selector.json")
+@@ -22,6 +22,16 @@ P2_METHODS = (
+ "bp", "fa", "dfa", "pepita", "ff", "ep", "dualprop", "clean_kp",
+ "sdil")
+ P2_ARCHITECTURES = ("minicnn", "vgglike", "vgg16")
++HARDWARE_PROFILES = {
++ "timan107_gtx1080": {
++ "allowed_physical_gpu_indices": [5, 7],
++ "allowed_device_names": ["NVIDIA GeForce GTX 1080"],
++ },
++ "a6000": {
++ "allowed_physical_gpu_indices": None,
++ "allowed_device_names": ["NVIDIA RTX A6000"],
++ },
++}
+
+
+ def output(*args):
+@@ -185,7 +195,7 @@ def p2_registry_sha256(jobs):
+ return hashlib.sha256(encoded).hexdigest()
+
+
+-def ensure_p2_launch(source, jobs):
++def ensure_p2_launch(source, jobs, hardware_policy):
+ path = os.path.join(ROOT, "runs", "plain-p2-launch.json")
+ immutable = {
+ "stage": "p2",
+@@ -194,6 +204,7 @@ def ensure_p2_launch(source, jobs):
+ "selector_sha256": SELECTOR_SHA256,
+ "num_registered_cells": len(jobs),
+ "registry_sha256": p2_registry_sha256(jobs),
++ "hardware_policy": hardware_policy,
+ "jobs": jobs,
+ }
+ if os.path.exists(path):
+@@ -413,11 +424,22 @@ def completed_histogram(experiment_name):
+ return paths[0]
+
+
+-def physical_gpu_report():
++def physical_gpu_report(profile):
+ index = os.environ.get("CUDA_VISIBLE_DEVICES")
+- if index not in ("5", "7"):
++ if index is None or not index.isdigit():
+ raise RuntimeError(
+- "formal P2 is authorized only on physical GPUs 5 and 7")
++ "formal P2 requires CUDA_VISIBLE_DEVICES to contain one "
++ "physical numeric GPU index")
++ if profile not in HARDWARE_PROFILES:
++ raise ValueError(f"unknown hardware profile: {profile}")
++ configured = HARDWARE_PROFILES[profile]
++ allowed_indices = configured["allowed_physical_gpu_indices"]
++ if (
++ allowed_indices is not None
++ and int(index) not in allowed_indices
++ ):
++ raise RuntimeError(
++ f"physical GPU {index} is not allowed by {profile}")
+ query = output(
+ "nvidia-smi",
+ "--query-gpu=index,uuid,name,memory.total",
+@@ -427,7 +449,13 @@ def physical_gpu_report():
+ if len(matches) != 1:
+ raise RuntimeError(f"cannot resolve physical GPU {index}")
+ physical_index, uuid, name, memory_mib = matches[0]
++ if name not in configured["allowed_device_names"]:
++ raise RuntimeError(
++ f"{profile} requires {configured['allowed_device_names']}, "
++ f"found {name}")
+ return {
++ "hardware_profile": profile,
++ "cuda_visible_devices": index,
+ "physical_index": int(physical_index),
+ "uuid": uuid,
+ "name": name,
+@@ -435,6 +463,19 @@ def physical_gpu_report():
+ }
+
+
++def hardware_policy(profile, gpu):
++ configured = HARDWARE_PROFILES[profile]
++ allowed_indices = configured["allowed_physical_gpu_indices"]
++ if allowed_indices is None:
++ allowed_indices = [gpu["physical_index"]]
++ return {
++ "profile": profile,
++ "allowed_physical_gpu_indices": list(allowed_indices),
++ "allowed_device_names": list(configured["allowed_device_names"]),
++ "single_gpu_per_process": True,
++ }
++
++
+ def descendant_pids(root_pid):
+ listing = subprocess.run(
+ ["ps", "-eo", "pid=,ppid="], check=True, capture_output=True,
+@@ -597,7 +638,7 @@ def terminate_process_group(process):
+ process.wait()
+
+
+-def run_p2_job(job, source, dry_run):
++def run_p2_job(job, source, gpu, dry_run):
+ existing = record_path(job["experiment_name"])
+ if existing is not None:
+ print(f"preserving {job['experiment_name']} -> {existing}", flush=True)
+@@ -609,7 +650,6 @@ def run_p2_job(job, source, dry_run):
+ print("RUN", " ".join(job["command"]), flush=True)
+ if dry_run:
+ return
+- gpu = physical_gpu_report()
+ started = time.time()
+ process = subprocess.Popen(
+ job["command"], cwd=ROOT, start_new_session=True)
+@@ -666,9 +706,9 @@ def run_p2_job(job, source, dry_run):
+ flush=True)
+
+
+-def run_job(job, source, dry_run):
++def run_job(job, source, gpu, dry_run):
+ if job["stage"] == "p2":
+- return run_p2_job(job, source, dry_run)
++ return run_p2_job(job, source, gpu, dry_run)
+ result = completed_histogram(job["experiment_name"])
+ if result is not None:
+ print(f"preserving {job['experiment_name']} -> {result}", flush=True)
+@@ -703,14 +743,23 @@ def main():
+ parser.add_argument("--shard-index", type=int, default=0)
+ parser.add_argument("--num-shards", type=int, default=1)
+ parser.add_argument("--method")
++ parser.add_argument(
++ "--hardware-profile",
++ choices=tuple(HARDWARE_PROFILES),
++ default=os.environ.get(
++ "SDIL_HARDWARE_PROFILE", "timan107_gtx1080"),
++ )
+ parser.add_argument("--dry-run", action="store_true")
+ args = parser.parse_args()
+ if not 0 <= args.shard_index < args.num_shards:
+ raise ValueError("invalid shard")
+ source = source_report()
+ jobs = p1_jobs() if args.stage == "p1" else p2_jobs()
++ gpu = None
+ if args.stage == "p2" and not args.dry_run:
+- launch = ensure_p2_launch(source, jobs)
++ gpu = physical_gpu_report(args.hardware_profile)
++ launch = ensure_p2_launch(
++ source, jobs, hardware_policy(args.hardware_profile, gpu))
+ print(f"P2 launch lock: {launch}", flush=True)
+ if args.method:
+ jobs = [job for job in jobs if job["method"] == args.method]
+@@ -721,7 +770,7 @@ def main():
+ if not jobs:
+ raise ValueError("no jobs selected")
+ for job in jobs:
+- run_job(job, source, args.dry_run)
++ run_job(job, source, gpu, args.dry_run)
+
+
+ if __name__ == "__main__":
+--
+2.54.0
+
diff --git a/sdil/data.py b/sdil/data.py
index 7f49b4e..759d7f7 100644
--- a/sdil/data.py
+++ b/sdil/data.py
@@ -2,13 +2,14 @@
so it runs on the Pascal env on the 1080 farm). Datasets already live on the
shared NFS home; the box is offline for dataset mirrors."""
+import os
import pickle
import struct
import hashlib
import numpy as np
import torch
-DATA_DIR = "/home/yurenh2/sdrn/data"
+DATA_DIR = os.environ.get("SDIL_DATA_DIR", "/home/yurenh2/sdrn/data")
# (subdir, mean, std) matching the standard torchvision normalisation used by
# the predecessor project, so numbers are comparable.