Fix JSEP pooling output shape to honor ceil_mode

[titaiwangms]

Summary

The ORT-Web JSEP pooling output-shape helper (PoolConvUtil.computePoolOutputShape → computeShapeHelper → adjustPadAndReturnShape) was floor-only: it had no ceilMode parameter and silently ignored ceil_mode. Under ceil_mode=1 this allocated the output tensor one element too small along each affected spatial axis, producing a wrong output shape (and downstream shape mismatches) — independent of any pooling divisor concern.

The same floor-only helper is duplicated in the legacy js/web/lib/onnxjs/util.ts; both copies are fixed.

Fix

• Thread a ceilMode parameter (default 0, i.e. floor) through computePoolOutputShape → computeShapeHelper → adjustPadAndReturnShape in both js/web/lib/wasm/jsep/util.ts and js/web/lib/onnxjs/util.ts.
• Route the NOTSET / VALID / SAME_UPPER / SAME_LOWER branches through a new computeOutputSize() helper that produces results identical to the C++ reference PoolAttributes::ComputeOutputSize (onnxruntime/core/providers/cpu/nn/pool_attributes.h), including the ceil_mode "shrink the last window if it starts entirely in the trailing padding" rule (ref: Fix corner case where output size need to reduce by one in MaxPool onnx/onnx#5741).
• The floor path (ceilMode default) is algebraically unchanged, so Conv and auto_pad pad-adjustment are unaffected.
• js/web/lib/wasm/jsep/webgpu/ops/pool.ts now passes attributes.ceilMode into the shape computation.

Scope: SHAPE-only

This PR fixes the output-shape computation only. End-to-end ceil_mode execution remains gated by the existing throw guards in parseAveragePoolAttributes / parseMaxPoolAttributes, because the WebGPU pooling kernel does not yet implement ceil_mode trailing-padding handling (and, for AveragePool, the count_include_pad divisor). Removing those throws + adding kernel support is a tracked follow-up; the now-correct shape path is exercised directly by the added unit tests until then. Comments at the throw sites and in the test header document this.

Tests

Adds js/web/test/unittests/pool-output-shape.ts (registered in unittests/index.ts), asserting output shapes for both implementations (jsep + onnxjs) against the C++ CPU reference ground truth:

• test_maxpool_2d_ceil → [1,1,2,2], AveragePool_10_ceil1_2d → [1,1,2,3]
• AvgPool ceil 1D [1,2,4] / 2D [1,1,3,3] / 3D [1,1,2,2,2] (matching the CPU pool_op_test.cc cases)
• ceil_mode with dilation, VALID / SAME_UPPER / SAME_LOWER auto_pad
• a shrink-rule discriminator (naive ceil() would give 3; correct = 2)
• floor-mode no-regression case

Related

• PyTorch context: ONNX export mismatch for avg_pool2d with ceil_mode=True and count_include_pad=True pytorch/pytorch#183528
• Related ORT PR: Align AvgPool ceil_mode on last value to torch #16752

Post-review fixups

• SAME_UPPER/SAME_LOWER integer division (external-review Major). The legacyTargetSize = (inSize + stride - 1) / stride computation now uses Math.floor(...) to match C++ pool_attributes.h ComputeSizePadDilations, which uses integer division. This fixes a latent float-division divergence that mis-rounded the auto_pad pad distribution. The correction applies to all ceil modes (not only ceil_mode=1), aligning JSEP/onnxjs with the CPU/CUDA EPs. The floor-path (ceil_mode=0) output shape is unchanged — only the SAME_* pad split is corrected — so nothing previously-correct regresses. Added ceil_mode=0 SAME_UPPER/SAME_LOWER non-divisible regression tests that assert the corrected pad out-param (SAME_UPPER → [0,1], SAME_LOWER → [1,0]).
• Throw-message / TODO clarity. The retained ceil_mode throw statements and the pool op TODO banner were reworded to make explicit that the output shape is now computed correctly, while ceil_mode kernel execution (padding/divisor) remains the pending WebGPU follow-up.

[Copilot]

Pull request overview

This PR fixes ORT-Web pooling output-shape computation to correctly honor ceil_mode (including the ONNX “shrink the last window if it starts entirely in trailing padding” rule) in both the JSEP and legacy onnxjs copies of PoolConvUtil.computePoolOutputShape. It also adds unit tests that exercise the pure shape helper directly (since end-to-end WebGPU pooling still throws on ceil_mode != 0 pending kernel support).

Changes:

• Thread ceilMode through computePoolOutputShape → computeShapeHelper → adjustPadAndReturnShape and centralize per-dimension sizing in a computeOutputSize() helper.
• Pass attributes.ceilMode into the WebGPU pooling output-shape computation.
• Add unit tests covering ceil-mode shapes (including the shrink rule) for both implementations (jsep + onnxjs).

Reviewed changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
js/web/lib/wasm/jsep/util.ts	Adds ceilMode plumbing and computeOutputSize() for pooling/conv output shape calculation.
js/web/lib/onnxjs/util.ts	Mirrors the same ceilMode and computeOutputSize() logic for the legacy onnxjs copy.
js/web/lib/wasm/jsep/webgpu/ops/pool.ts	Passes ceilMode into shape computation and expands comments around current ceil_mode execution guards.
js/web/test/unittests/pool-output-shape.ts	New unit tests validating ceil-mode output shapes (and shrink rule) against CPU-reference expectations.
js/web/test/unittests/index.ts	Registers the new unit test module.

[titaiwangms]

Review summary

Solid, well-tested, and correctly scoped (shape-only, execution still gated). The ceil-mode math matches the C++ reference (PoolAttributes::ComputeOutputSize), all test expectations check out (including the shrink-rule boundary tested exactly at equality), the two computeOutputSize copies are byte-identical, and all other callers (Conv, WebGL, GlobalPool) safely keep the floor default. One Major to address before the ceil path can be fully trusted; the rest are Minor/Nit.

🔴 Major (1)

SAME_UPPER/SAME_LOWER uses float division for legacyTargetSize, diverging from the C++ reference in ceil mode.
js/web/lib/wasm/jsep/util.ts:523 and js/web/lib/onnxjs/util.ts:1420

const legacyTargetSize = (inSize + stride - 1) / stride;   // float; C++ uses int64 division

Reproduced with inSize=2, stride=2, kernel=3, SAME_UPPER, ceilMode=1:

• JS: legacyTargetSize = 1.5 → padNeeded = 2 → output […,2]
• C++ (legacy_target_size = 1) → pad_needed = 1 → output […,1]

Floor mode masks this (the extra pad is absorbed by floor), so it's a pre-existing latent bug that the new ceil path newly exposes — and it breaks the "identical to the C++ reference" claim for the SAME_* branch. The added SAME_* tests only use in=5, stride=2, where the division is integral and hides the issue.

Fix: use integer division to match C++, in both files:

const legacyTargetSize = Math.floor((inSize + stride - 1) / stride);

and add a non-divisible regression case, e.g. inputDims=[1,1,2], strides=[2], kernelShape=[3], autoPad='SAME_UPPER', ceilMode=1 → expected [1,1,1].

🟡 Minor (2)

1. float32 vs float64 ceil. C++ does std::ceil(static_cast<float>(numerator) / stride) (float32); JS uses double Math.ceil (jsep/util.ts:488, onnxjs/util.ts:1385). These diverge only for numerator > 2^24, which is unreachable for real pooling spatial dims — so no practical impact. Either emulate float32 with Math.fround(...) for strict parity, or soften the "identical to the C++ reference" comment to note JS intentionally uses exact double math. Non-blocking.
2. Stale banner TODO at the top of pool.ts still lists ceil_mode as an unqualified TODO, which now contradicts the precise per-function comments in parseAveragePoolAttributes/parseMaxPoolAttributes. Suggest updating it to note the output-shape math is supported and only kernel execution is gated.

⚪ Nits

• The "Keep in sync" comment on the duplicated computeOutputSize isn't enforced by anything but hand-picked case parity. Since the test already imports both PoolConvUtils, consider an explicit "both implementations agree" test over a matrix of random inputs — turns the promise into CI coverage.
• computeOutputSize's numerator param is semantically indirect; and the pool-output-shape.ts "Reviewer-verified" comment on the SAME_UPPER case would be more traceable as the actual hand-derivation.

Follow-up (out of scope, already tracked)

End-to-end ceil_mode execution remains correctly gated by the throws in parseAveragePool/MaxPoolAttributes pending the kernel trailing-padding (and AveragePool count_include_pad) support — the tracked follow-up.

Reviewed by a multi-model review team (readability / correctness / adversarial / spec-adherence / cross-module integration).

[titaiwangms]

Review fold-in (post-review fixups pushed: 134ab29f8, 1ea30f903)

Thanks all for the thorough passes — and a special thanks to the external review team for catching the SAME_* integer-division issue.

Folded in ✅

• SAME_UPPER/SAME_LOWER integer-division Major (external team). legacyTargetSize now uses Math.floor((inSize + stride - 1) / stride) to match C++ pool_attributes.h ComputeSizePadDilations (integer division). This was a latent float-division divergence that mis-rounded the auto_pad pad split. It's live on all ceil modes — on the ceil_mode=0 floor path the output shape is unchanged but the pad distribution was wrong (e.g. in=2, stride=2, kernel=3, SAME_UPPER gave pads [1,1] instead of C++ [0,1]). Added ceil_mode=0 SAME_UPPER/SAME_LOWER non-divisible pad-distribution regression tests (assert the corrected out-param, not just the shape); they fail against the pre-fix float code and pass with the fix.
• Throw-message / TODO clarity (Copilot C3/C4 + team Minor 2). Reworded the retained ceil_mode throw messages and the pool op TODO banner to make clear the output shape is computed, while ceil_mode kernel execution (padding/divisor) is the pending WebGPU follow-up.

Respectfully declined (with rationale) 🙏

• Copilot C1/C2 — Math.floor → Math.trunc on the floor path. These differ only for a negative numerator with a remainder, i.e. kernel > inSize + pads producing an invalid <= 0 output dim — unreachable for valid models. It's also the ceil_mode=0 path this PR intentionally scopes out. Happy to file separately if desired.
• Team Minor 1 — float32 vs double for the ceil computation. Diverges only for numerator > 2^24, impossible for real spatial dims; non-blocking, as the team itself noted.
• Team Nits A/B — random-matrix cross-impl test + numerator rename. Optional polish; left out to keep the fixup tight and reviewable.

[hariharans29]

Review: PR #29627 — Fix JSEP pooling output shape to honor ceil_mode

Approve. Well-scoped correctness fix. PoolConvUtil.computePoolOutputShape in both js/web/lib/wasm/jsep/util.ts (JSEP) and js/web/lib/onnxjs/util.ts (legacy onnxjs) previously ignored ceil_mode, so under ceil_mode=1 the output tensor was allocated one element too small on affected spatial axes. This PR threads a ceilMode parameter through computePoolOutputShape → computeShapeHelper → adjustPadAndReturnShape and centralizes per-dimension sizing in a new computeOutputSize() helper whose body is a line-for-line port of the C++ PoolAttributes::ComputeOutputSize. Scope stays SHAPE-only — the throw guards in parseAveragePoolAttributes/parseMaxPoolAttributes remain, with reworded messages that make the split (shape vs kernel-execution) explicit. CI is green at 86/86.

Correctness — matches the C++ reference

Cross-checked against onnxruntime/core/providers/cpu/nn/pool_attributes.h ComputeOutputSize and ComputeSizePadDilations:

1. Numerator formula. JS numerator = inSize + padHead + padTail - dkernel with dkernel = dilation * (kernel - 1) + 1 → same as C++ int64_t numerator = in_size + pad_head + pad_tail - dilation * (kernel - 1) - 1.
2. Floor branch. JS Math.floor(numerator / stride) + 1 matches C++ numerator / stride + 1 on the reachable (non-negative-numerator) domain. The Math.floor vs Math.trunc divergence on negative numerators is unreachable for valid models (kernel <= inSize + pads is a model-validity precondition), and Copilot C1/C2 was correctly declined for that reason.
3. Ceil branch. JS Math.ceil(numerator / stride) + 1 matches C++ static_cast<int64_t>(std::ceil(static_cast<float>(numerator) / stride)) + 1. The float32-vs-double divergence only shows up for numerator > 2^24, unreachable for real spatial dims.
4. Shrink rule. JS if ((outSize - 1) * stride >= inSize + padHead) outSize -= 1 is byte-identical to C++ if ((out_size - 1) * stride >= in_size + pad_head) --out_size (both signed) and correctly cites the ONNX rationale (onnx/onnx#5741).
5. SAME_UPPER/SAME_LOWER integer division. The external-team Major (legacyTargetSize was using float division) is now fixed with Math.floor((inSize + stride - 1) / stride). Nice catch — this was a pre-existing latent bug on the floor path too: the output shape happened to be right because Math.floor(...+ 1) absorbed the extra half-unit, but the pad distribution written back to the pads out-param was wrong (in=2, stride=2, kernel=3, SAME_UPPER used to give [1,1] instead of the correct [0,1]). Downstream consumers of that pads array (e.g. anything that actually places the sliding window relative to pad_head) were silently getting the wrong window offset. The PR adds regression tests that assert on expectedPads for both SAME_UPPER→[0,1] and SAME_LOWER→[1,0] at ceil_mode=0, which is exactly the right guard.

I hand-verified each new test case against the C++ formula. All expected shapes match. Spot check of the shrink-rule discriminator: inSize=3, stride=2, kernel=2, pads=[1,1], ceil_mode=1:

• numerator = 3 + 1 + 1 − 1 − 1 = 3, ceil(3/2) + 1 = 3, shrink test (3−1)*2 = 4 >= 3+1 = 4 → true → shrink → 2. ✓ (matches expectedShape: [1,1,2]).

The shrink test is exercised exactly at equality (4 >= 4), which is the boundary that would be the most likely to slip in a signed/unsigned mismatch. Good.

Scoping

1. Backwards-compatible defaults. ceilMode = 0 on all touched signatures. Conv, GlobalPool (via globalPoolAttributes which stays at ceilMode: 0), the onnxjs WebGL caller (which intentionally does not thread ceilMode), and any pre-existing pool callers all continue to get the pre-PR floor semantics. No legitimate caller regresses.
2. Execution still gated. The throws in parseAveragePoolAttributes and parseMaxPoolAttributes remain and are reworded to reflect the split ("output shape is computed, but kernel execution … is not yet implemented"). The pool.ts banner TODO is updated to match. This is the right containment — the WebGPU kernel still needs trailing-padding handling (and for AveragePool, count_include_pad divisor handling) before ceil_mode end-to-end works, and shipping a correct shape helper while gating execution keeps the fix reviewable and unblocks whatever caller was already tripping the shape assertion.
3. Duplicated computeOutputSize in onnxjs + jsep. The two implementations are line-for-line identical and the test matrix runs both against the same expected values, so silent drift would fail CI. The "keep in sync" comment is not compile-time-enforced, but the test coverage is the effective enforcement. Acceptable.

Test quality

1. Sixteen shape cases against the C++ oracle, each with a comment tying it back to a specific CPU test case or C++ formula. test_maxpool_2d_ceil, AveragePool_10_ceil1_2d, the AvgPool 1D/2D/3D ceil family from pool_op_test.cc, dilation, VALID / SAME_UPPER / SAME_LOWER auto_pad, shrink-rule boundary, floor-path no-regression, and the SAME_* integer-division regression guards. That's an honest coverage matrix.
2. Every case runs against both PoolConvUtil copies (jsep + onnxjs), so cross-copy drift is caught.
3. runPoolConvUtil slices strides/dilations/kernelShape/pads before calling computePoolOutputShape — correctly defends against mutation of shared inputs across the two implementations (adjustPadAndReturnShape writes back to pads in the auto_pad branches).
4. The expectedPads assertion on the two ceil_mode=0 SAME_* non-divisible cases is exactly the right thing to catch the latent SAME_* pad-distribution bug — asserting only shape would have missed it because the pre-fix output shape was accidentally correct.
5. defaults to floor when ceilMode is omitted guards the default parameter contract.

Non-blocking observations

1. Cross-copy parity test. Team Nit A (a randomized cross-implementation test) would harden the "keep in sync" comment into CI enforcement. The hand-picked matrix is already good; a small fast-check-style randomized case would make silent divergence structurally impossible. Fine to leave for a follow-up.
2. computeOutputSize param shape. Team Nit B (numerator param feels indirect; callers still separately pass inSize and padHead for the shrink test). Alternative would be computeOutputSize(inSize, padHead, padTail, kernel, stride, dilation, ceilMode) and let the helper build numerator internally — a hair more overhead but a cleaner interface. Cosmetic.
3. Math.trunc vs Math.floor on the floor path. Author correctly declined this as unreachable for valid models. If anyone ever wants absolute-strict C++ parity even in the malformed-input case, Math.trunc would be the exact port; the current choice is defensible. A one-line comment noting "matches C++ integer division on the non-negative-numerator domain" wouldn't hurt.
4. Legacy onnxjs coverage. The PR touches js/web/lib/onnxjs/util.ts and the internal comment there notes "the onnxjs WebGL pooling caller intentionally does NOT pass ceilMode (legacy path still throws on ceil_mode != 0), so it always uses the floor default." Worth spot-checking that backends/webgl/ops/pool.ts still throws on ceilMode != 0, i.e. this PR doesn't accidentally unlock a broken WebGL execution path. From what I can see the WebGL parse path is untouched by this PR, so the throw is preserved — good — but a one-line note in the PR description confirming that would help future readers.
5. legacyTargetSize fix scope. The Math.floor correction to legacyTargetSize corrects the pad distribution on the ceil_mode=0 floor path too (unrelated to ceil_mode). That's genuinely a correctness fix for pre-existing behavior, not just a ceil-mode enabler. Worth calling out more prominently in the commit message / squash-merge message so it's discoverable if a regression is later suspected somewhere unrelated to ceil_mode.
6. Auto-merge is already enabled — that's fine given CI is green, but if hariharans29's review lands changes, the auto-merge will consume the follow-up commits. Not a concern, just an FYI.

Summary

Approve. Fix is correct against the C++ reference (verified by walking each test case through the formulas), correctly scoped as shape-only with clear kernel-execution follow-up documented in place, the external-team-caught SAME_* integer-division bug is fixed with a proper regression test that specifically pins the pad out-param, and both jsep and onnxjs copies stay in lockstep through the shared test matrix. The declined polish items are appropriately declined. Good to go once @hariharans29 weighs in.