[Panadestein]

🤖 AI text below 🤖

What

• SiteKernels plans every contraction with opt_einsum's dp optimizer instead of the default auto heuristic.
• peak_elements walks the same path, so the planner sizes batches for the path the kernel runs. Planner policy and batch rounding are unchanged.
• src, apply and compress take optimize (default "dp"): any path optimizer name opt_einsum resolves, including ones added with register_path_fn, or a PathOptimizer instance. An unknown name raises ValueError, and anything else raises TypeError. Documented in Large problems → Contraction paths and the README.
• benches/large/bench_einsum_dp.py, with the results in benches/large/README.md (no raw results or job scripts, like the other benches).

Why: measured on A100-SXM-64GB, complex128

Full sweep, D_M = 1000, chi_out = 500, 60 GB budget, 3 repeats:

GPUs	auto median (range)	dp median (range)	speed-up	QR distance
1	109.7 s (109.6-130.6)	56.2 s (56.2-74.1)	1.95x	2.7e-13
4	32.4 s (32.3-51.7)	17.4 s (17.4-40.1)	1.86x	4.0e-13

• The planner chose the same batches with both optimizers. dp's planned peaks are lower: 50.3 against 58.8 GB on one GPU, and 12.2 against 32.3 GB per GPU on four.
• All of the gain is in the right-to-left pass (95.4 → 42.3 s on one GPU). At these shapes auto's sketch costs 3.9x dp's flops, and its projection 2.0x.

Kernels at the reference shapes (D = (4, 4, 4000, 4), eta = 2000):

• Projection: dp halves the time per row (25.2 → 11.8 ms at 8 rows, 22.1 → 11.0 ms at 16). From 32 rows up, auto plans an 8.2e9-element intermediate (265-283 GB), which holds the planner to 16 rows. dp runs up to 128 rows in 10-25 GB.
• Sketch: dp is 3.7-3.8% slower. Its path has only 0.5% fewer flops.
• Environment step: identical paths.
• Agreement: all outputs agree to 1.6e-14 or better.
• Path search: about 1 ms per shape for dp. Paths are cached per src call, so a sweep pays 0.03-0.10 s.

Not measured: a reference-size sweep. The kernel times reproduce its old right-to-left pass (3140 s estimated, 3112 s measured), and predict about 1.5x with dp.

Testing

• pytest -m "not slow": 370 passed and 6 skipped (multi-GPU tests, on one GPU), on a Leonardo A100. The latest run, with optimize as an argument, gave 387 passed.
• New tests check that auto, greedy, a DynamicProgramming instance and a registered name give the same result, that invalid values raise through all three entry points, and that peak_elements follows the optimizer.
• prek run --all-files passes.

Notes for review

• This changes the default contraction paths for every user. Outputs agree to rounding, not bit for bit.

🤖 Generated with Claude Code

[github-actions]

Test Results

    5 files  ±  0      5 suites  ±0   3m 55s ⏱️ -21s
  394 tests + 20    387 ✅ + 20   7 💤 ±0  0 ❌ ±0 
1 890 runs  +100  1 875 ✅ +100  15 💤 ±0  0 ❌ ±0 

Results for commit 507b768. ± Comparison against base commit bd78e23.

♻️ This comment has been updated with latest results.

[robertodr]

LGTM, wouldn't it be simpler to just add it as an option exposed through the API though (instead of an env-var)? This way whatever new method becomes available in opt-einsum, it could be immediately used, without having to update the PathOptimizer list.

[Panadestein]

🤖 AI text below 🤖

@robertodr Agreed, done in 507b768. The env var is gone. src, apply and compress take optimize (default "dp") and pass it to the kernels and to the memory plan, so batches are still sized for the path that runs.

It accepts anything opt_einsum can plan every kernel with:

• any name opt_einsum resolves, checked with opt_einsum.paths.get_path_fn, so new built-ins and names added with register_path_fn work without a list on our side;
• a PathOptimizer instance, e.g. oe.DynamicProgramming(minimize="size").

Booleans and explicit paths are refused with a TypeError, because one path cannot fit the sweep's different equations. The fast suite passes on an A100 (387 passed).