[#449 inc4] Container-type closeout: SequenceMap + cross-subgraph capture + Scan state

[justinchuby]

#449 increment 4 — closeout slice (DRAFT — do NOT merge)

Rebased onto main now that inc3a (#527) is merged — a trivial superset (single commit on top of main). CPU-only, pure onnx-runtime-shape-inference crate.

Closes out the #449 container-type arc: foundation → seq ops → seq↔tensor → If/Loop control flow → this.

What landed

1. SequenceMap (opset 17) — maps a subgraph body over sequence element(s), wrapping each body output back as a Sequence. Handled in infer_control_flow like If/Loop/Scan (subgraph ops are not registry rules). Zips multiple input sequences; broadcasts extra tensor inputs. Composed from the inc3a body-threading driver (set_body_input / body_container_seeds / read_body_output / map_container_to_parent) — no fork.
2. Cross-subgraph container capture — remap_node_io previously dropped a captured value's ValueType (value_type: None). New remap_container_type threads the full container type (recursing, remapping tensor-leaf dims via remap_dim_expr, mirroring the tensor path). bind_captures / extend_visible_scope gained a containers channel, so a body referencing an outer-scope sequence now sees its element type.
3. Scan container state (folded in — previously deferred inc3b) — Scan state vars can be containers; reuses seed_container_input / body_container_seeds + infer_scan_outputs' container check. All four subgraph ops now container-complete.

Hard bar

• Byte-identical tensor path preserved — every new path is gated on a container being present; regression test tensor_only_path_is_byte_identical_after_container_type_model stays GREEN.
• DRY — reuses inc3a body-threading helpers + foundation unify_value_type/unify_tensor_type/merge_element_shape. No per-op copy-paste.
• Catalog stays 217 ops / 262 entries (HONEST) — SequenceMap is a control-flow op, not a registry rule (like If/Loop/Scan); there is no 218th handler, so bumping the pin would be a phantom declaration.
• Tests: graph_inference 33 → 39 — SequenceMap identity/transform/zip/broadcast, If capturing an outer-scope sequence (proves the remap_node_io fix), Scan carrying a sequence state var.

Verification (all GREEN, post-rebase)

cargo test -p onnx-runtime-shape-inference (16 op_rules + 39 graph_inference + 275 unit + 1 doctest); cargo fmt --all --check; cargo clippy -p onnx-runtime-shape-inference --all-targets -- -D warnings; cargo build -p onnx-runtime-session -p onnx-runtime-eager.

Recommendation

Close #449. The sequence container-type arc is complete end to end. Only Optional/Map op handlers remain — not implemented, not used by any in-tree target model, no real-model demand → optional low-priority follow-up, not a load-bearing gap.

refs #449

[codecov]

Codecov Report

❌ Patch coverage is 79.36508% with 39 lines in your changes missing coverage. Please review.
✅ Project coverage is 80.59%. Comparing base (5a36f9d) to head (f99fd65).

Files with missing lines	Patch %	Lines
crates/onnx-runtime-shape-inference/src/infer.rs	79.36%	33 Missing and 6 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #531      +/-   ##
==========================================
- Coverage   80.59%   80.59%   -0.01%     
==========================================
  Files         315      315              
  Lines      123096   123259     +163     
  Branches   123096   123259     +163     
==========================================
+ Hits        99209    99337     +128     
- Misses      19859    19890      +31     
- Partials     4028     4032       +4     

Flag	Coverage Δ
cli-ort-linux	83.27% <ø> (ø)
cli-ort-windows	82.67% <ø> (-0.11%)	⬇️
mlas	77.91% <ø> (ø)
offline	80.51% <79.36%> (-0.01%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines	Coverage Δ
crates/onnx-runtime-shape-inference/src/infer.rs	90.04% <79.36%> (-2.00%)	⬇️

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[github-actions]

🔴 Benchmark Regression Detected

Comparison of criterion micro-benchmarks: PR head vs merge-base, measured on the same runner in the same job (base first → PR second).

ℹ️ Absolute times are informational only — they vary with runner load. The % change column is the reliable signal because both sides ran under identical conditions.

Status	Scenario	Base	PR	Change
🔴	matmul/medium_generic_f32_threads=8/32x512x512	1.04 ms	1.40 ms	+35.2%
⚠️	logit_processing/seven_processor_chain_per_step	1.34 ms	1.69 ms	+26.1%
⚠️	kv_cache/alloc_dealloc_pages	41.54 µs	50.83 µs	+22.4%
⚠️	tokenization/decode_tokens_per_second	6.48 ms	7.59 ms	+17.2%
✅	add/small_bf16_threads=1-internal/1024	12.90 µs	14.13 µs	+9.5%
✅	add/small_f32_threads=1-internal/1024	200.6 ns	216.0 ns	+7.7%
✅	gather/large_f32_threads=1-internal/131072	41.77 µs	43.47 µs	+4.1%
✅	matmul/small_generic_bf16_threads=1/1x256x256	37.90 µs	39.34 µs	+3.8%
✅	tokenization/encode_tokens_per_second	391.22 µs	401.30 µs	+2.6%
✅	matmul/small_generic_f16_threads=8/1x256x256	38.42 µs	39.23 µs	+2.1%
✅	gather/large_f16_threads=1-internal/131072	12.90 µs	13.14 µs	+1.8%
✅	add/small_f16_threads=1-internal/1024	15.12 µs	15.16 µs	+0.3%
✅	add/medium_bf16_threads=1-internal/262144	2.64 ms	2.64 ms	+0.1%
✅	add/large_f16_threads=1-internal/4194304	41.12 ms	40.99 ms	-0.3%
✅	add/medium_f32_threads=1-internal/262144	2.68 ms	2.67 ms	-0.3%
✅	add/medium_f16_threads=1-internal/262144	2.60 ms	2.59 ms	-0.5%
✅	grammar_masking/llguidance_compute_mask/32	88.09 µs	87.21 µs	-1.0%
✅	reduce_mean/small_f32_threads=1-internal/4096	15.26 µs	15.08 µs	-1.2%
✅	matmul/large_generic_bf16_threads=8/32x1024x1024	1.75 ms	1.71 ms	-2.1%
✅	sampling_latency/min_p_per_token	437.53 µs	428.32 µs	-2.1%
✅	add/large_f32_threads=1-internal/4194304	41.69 ms	40.69 ms	-2.4%
✅	matmul/medium_generic_bf16_threads=1/32x512x512	579.80 µs	565.12 µs	-2.5%
✅	matmul/medium_generic_f16_threads=8/32x512x512	37.57 µs	35.73 µs	-4.9%
✅	matmul/medium_generic_f32_threads=1/32x512x512	2.65 ms	2.52 ms	-5.0%
✅	matmul/small_generic_bf16_threads=8/1x256x256	45.48 µs	42.80 µs	-5.9%
✅	gather/small_f32_threads=1-internal/4096	770.7 ns	724.6 ns	-6.0%
✅	gather/small_f16_threads=1-internal/4096	502.1 ns	469.1 ns	-6.6%
✅	add/large_bf16_threads=1-internal/4194304	45.55 ms	42.22 ms	-7.3%
✅	reduce_mean/medium_f32_threads=1-internal/65536	262.06 µs	242.58 µs	-7.4%
✅	sampling_latency/greedy_per_token	3.95 µs	3.51 µs	-11.2%
✅	matmul/large_generic_f16_threads=8/32x1024x1024	114.47 µs	101.46 µs	-11.4%
✅	reduce_mean/large_f32_threads=1-internal/262144	1.14 ms	994.11 µs	-12.9%
✅	gather/medium_bf16_threads=1-internal/32768	2.89 µs	2.50 µs	-13.5%
✅	gather/small_bf16_threads=1-internal/4096	566.1 ns	486.3 ns	-14.1%
✅	matmul/medium_generic_bf16_threads=8/32x512x512	506.35 µs	433.00 µs	-14.5%
✅	matmul/large_generic_bf16_threads=1/32x1024x1024	2.44 ms	2.09 ms	-14.5%
🟢	sampling_latency/top_p_per_token	1.45 ms	1.23 ms	-15.5%
🟢	matmul/medium_generic_f16_threads=1/32x512x512	38.33 µs	32.16 µs	-16.1%
🟢	gather/medium_f16_threads=1-internal/32768	2.94 µs	2.45 µs	-16.5%
🟢	matmul/small_generic_f16_threads=1/1x256x256	47.73 µs	39.29 µs	-17.7%
🟢	matmul/large_generic_f32_threads=1/32x1024x1024	12.61 ms	9.69 ms	-23.2%
🟢	gather/medium_f32_threads=1-internal/32768	5.08 µs	3.88 µs	-23.8%
🟢	matmul/large_generic_f16_threads=1/32x1024x1024	111.74 µs	84.95 µs	-24.0%
🟢	sampling_latency/top_k_per_token	633.09 µs	473.97 µs	-25.1%
🟢	matmul/small_generic_f32_threads=1/1x256x256	51.18 µs	37.90 µs	-26.0%
🟢	matmul/large_generic_f32_threads=8/32x1024x1024	6.19 ms	4.42 ms	-28.6%
🟢	gather/large_bf16_threads=1-internal/131072	16.79 µs	11.33 µs	-32.5%
🟢	matmul/small_generic_f32_threads=8/1x256x256	67.37 µs	34.50 µs	-48.8%

Visual flags: ⚠️ ≥ 15% slower, 🔴 ≥ 30% slower — calibrated against measured runner noise (~27% worst-case on multi-threaded matmul)

Host info

CPU: Apple M1 (Virtual)
Cores: 3
OS: Darwin 25.5.0 arm64
Rust: rustc 1.97.1 (8bab26f4f 2026-07-14)
Load avg: { 4.64 4.25 5.83 }

What this cannot catch

• Regressions in code paths not covered by these benchmarks (e.g., end-to-end decode with a real model)
• Sub-threshold regressions that compound over multiple PRs
• Performance changes that only manifest under GPU execution
• Latency changes in the ORT integration path (these benchmarks exercise the native Rust kernels)

[justinchuby]

VERDICT: APPROVE

Independent review of PR #531 (squad/449-container-inc4, f99fd65, base main, refs #449). CPU-only shape-inference crate. Author Harry locked out; I verified everything from source + a clean worktree build.

Verification (all GREEN)

• cargo test -p onnx-runtime-shape-inference: lib 16; graph_inference 39 (33→39, +6); op_rules 275; doctest 1. 0 failed. Counts match Harry's claim exactly.
• cargo fmt --all --check: clean (exit 0).
• cargo clippy -p onnx-runtime-shape-inference --all-targets -- -D warnings: clean (exit 0).
• Blast radius cargo build -p onnx-runtime-session -p onnx-runtime-eager: clean (exit 0). Container ValueType default-None does not break session/eager.

1. Byte-identical tensor path — CONFIRMED, gate is real for the new paths

Every new container path is gated on a container actually being present:

• remap_container_type runs only under io.value_type.as_ref().map(...) (infer.rs:1368-1371); tensor edges have value_type: None, so it is never entered.
• Scan container output gated on slot < num_state && report.containers.get(&body_output) (infer.rs:1138-1140).
• body_container_seeds Scan/Loop branches enter only via containers.get(&src) in seed_container_input (infer.rs:1015-1027); empty container map ⇒ no seeds.
• bind_captures inserts into containers only when binding.value_type is Some (infer.rs:1290-1294).
• SequenceMap seed/infer keyed on op_type == "SequenceMap", absent from tensor graphs.
  tensor_only_path_is_byte_identical_after_container_type_model (graph_inference.rs:1497) GREEN, plus the full 275-test op_rules regression + existing If/Loop/Scan tensor tests. One behavioral note (non-blocking): extend_visible_scope now publishes a scope binding when only shape_data is present (previously required a types entry) — this is monotonic (more resolution info, never wrong info), only affects subgraph-capture scope, and is covered by the fully-green suite.

2. SequenceMap semantics — CORRECT, general, DRY

• Body input i seeded from operand i's per-element type: sequence operand ⇒ element tensor via as_sequence_element().as_tensor().to_type_info(); plain tensor operand ⇒ whole tensor via types.get fallback (infer.rs:648-664). Matches ONNX opset-17 spec (zip input sequences; broadcast additional tensor inputs). Seq-of-seq element handled as a container seed (infer.rs:706-720).
• Each body output wrapped one level in Sequence (infer_sequence_map_outputs, infer.rs:1191-1234), including the seq-of-seq output case.
• DRY: reuses set_body_input, seed_container_input, body_container_seeds, read_body_output, map_container_to_parent — the same inc3a driver; body inference is not forked (only the output-wrapping wrapper is SequenceMap-specific, as it must be).
• Tests assert the OUTPUT sequence element type (not just non-crash): identity preserves f32 [2,3]; Shape transform ⇒ i64 [2]; zip of two seqs (Add) ⇒ f32 [2,3]; additional-tensor broadcast (seq + bias) ⇒ f32 [2,3] (graph_inference.rs:1877-2062).

3. Cross-subgraph capture fix — SOUND, both branches, not degraded

remap_container_type (infer.rs:1380-1418) recurses the container and remaps every tensor-leaf dim through the same remap_dim_expr + shared parent_to_child/child_to_parent maps used by the tensor type_info path in the same remap_node_io call. remap_dim_expr uses entry().or_insert_with (infer.rs:1432-1436), so one child symbol per parent symbol — no aliasing/collision, correct mirror of map_container_to_parent. if_branches_capture_outer_scope_sequence (graph_inference.rs:1968) builds seq0 in the outer scope, both If branches return it by lexical capture, and asserts the unified output is Sequence<f32 [4,5]> — a concrete container only resolvable if BOTH branches threaded the element type (a degrade-to-None in either would break unification/container output).

4. Scan container fold — CORRECT

State slots (slot < num_state) may carry containers; scan-output slots always stack tensors and are never wrapped (infer.rs:1136-1144). scan_carries_sequence_state_variable (graph_inference.rs:2065) asserts real state-carry: SequenceErase on the carried sequence still yields final_state = Sequence<f32 [2,3]>, AND asserts !report.containers.contains_key(&scan_result) — proving the stacked scan output stays a plain tensor.

5. Catalog honesty — CONFIRMED

217/262 unchanged is correct. SequenceMap is control-flow, handled in infer_control_flow (infer.rs:370), never registered as a registry op (grep confirms no registry entry). A +1 would be a phantom op failing the pin. expanded_registry_catalog_count_is_pinned (op_rules.rs:283-288) GREEN: operator_count 217, entry_count 262.

6. Close-#449 judgment (advisory)

Support closing #449. The container-type shape-inference arc is genuinely complete for the real-model surface: full Sequence family + threading through all four subgraph ops (If/Loop/Scan/SequenceMap) + cross-subgraph capture. Remaining Optional/Map op handlers have zero in-tree real-model demand (onnx-genai targets are tensor transformer/LLM graphs), and the ValueType::Optional/Map representation already exists for a small later follow-up. One known, non-blocking gap to record on the epic: container ValueType lives only in the parallel report.containers map and is not persisted onto the IR Value (no IR-level container-type persistence). That is invisible to downstream session/eager today (which is fine — the blast-radius build is green and those paths are tensor-only), but if a future consumer needs container types post-inference, IR persistence would be the next increment. Not a blocker for closing #449.

Recommendation: APPROVE and close the #449 epic.

— Lori (independent review)