fix: correct conv/pool padding, reshape 0/-1, reject unsupported ops at load

.github/workflows/ci.yml

@@ -3,8 +3,8 @@ name: CI
 on:
   push:
     branches: [main]
+  # all PRs, so stacked PRs (based on another branch) get checked too
   pull_request:
-    branches: [main]
 
 env:
   CARGO_TERM_COLOR: always
@@ -81,3 +81,6 @@ jobs:
 
       - name: Run ORT comparison tests
         run: .venv/bin/python tests/compare_ort.py
+
+      - name: Run op edge-case tests
+        run: .venv/bin/python tests/test_ops.py

CHANGELOG.md

@@ -7,7 +7,18 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Fixed
+- Conv2d / MaxPool2d / AveragePool2d with asymmetric pads no longer average the two
+  sides; uneven pads are applied to the input explicitly (exact for `count_include_pad=0`).
+- `auto_pad=SAME_UPPER` / `SAME_LOWER` / `VALID` are now honoured instead of treated as no padding.
+- Pooling now respects `ceil_mode`.
+- Reshape with both `0` and `-1` in the target shape no longer panics (divide by zero),
+  and the opset < 5 `shape` attribute form is supported.
+- Unsupported ops now fail at `load_onnx` with a list of the missing ops, instead of
+  being skipped with a warning and failing (or silently misbehaving) at inference time.
+
 ### Added
+- `tests/test_ops.py`: per-op edge-case tests against ONNX Runtime, run in CI.
 - ONNX interpreter support for Conv2d, fused BatchNormalization, MaxPool2d,
   AveragePool2d, and GlobalAveragePool.
 - GitHub Actions CI: `cargo build`, `cargo fmt` / `cargo clippy` (non-blocking),

README.md

@@ -41,6 +41,7 @@ maturin develop --release
 python tests/make_test_model.py   # generates tests/mlp.onnx
 python tests/test_bridge.py       # numpy <-> Burn tensor bridge
 python tests/compare_ort.py       # correctness + perf vs ONNX Runtime
+python tests/test_ops.py          # op edge cases (padding, pooling, reshape) vs ORT
 ```
 
 ## License

src/interpreter/mod.rs

@@ -138,6 +138,22 @@ pub fn load_onnx(path: &str) -> Result<OnnxModel, String> {
         .parse_file(path)
         .map_err(|e| e.to_string())?;
 
+    // fail fast: a skipped op would only surface later as a missing tensor or wrong output
+    let mut unsupported: Vec<String> = graph
+        .nodes
+        .iter()
+        .filter(|n| !is_supported(n))
+        .map(op_type)
+        .collect();
+    if !unsupported.is_empty() {
+        unsupported.sort();
+        unsupported.dedup();
+        return Err(format!(
+            "unsupported ONNX op(s): {}",
+            unsupported.join(", ")
+        ));
+    }
+
     let mut weights = HashMap::new();
 
     // pass 1: load all named static weights
@@ -167,6 +183,40 @@ pub fn load_onnx(path: &str) -> Result<OnnxModel, String> {
     })
 }
 
+// keep in sync with `dispatch`
+fn is_supported(node: &Node) -> bool {
+    matches!(
+        node,
+        Node::Relu(_)
+            | Node::Sigmoid(_)
+            | Node::Tanh(_)
+            | Node::Gelu(_)
+            | Node::Softmax(_)
+            | Node::LogSoftmax(_)
+            | Node::Linear(_)
+            | Node::Gemm(_)
+            | Node::Add(_)
+            | Node::Sub(_)
+            | Node::Mul(_)
+            | Node::Div(_)
+            | Node::Reshape(_)
+            | Node::Flatten(_)
+            | Node::Transpose(_)
+            | Node::Conv2d(_)
+            | Node::BatchNormalization(_)
+            | Node::MaxPool2d(_)
+            | Node::AveragePool2d(_)
+            | Node::GlobalAveragePool(_)
+            | Node::Constant(_)
+    )
+}
+
+/// Op type of a node (the enum variant name, e.g. "MatMul"), for error messages.
+fn op_type(node: &Node) -> String {
+    let dbg = format!("{node:?}");
+    dbg.split('(').next().unwrap_or(&dbg).to_string()
+}
+
 fn dispatch(node: &Node, ctx: &mut ExecutionContext) {
     match node {
         Node::Relu(n) => ops::activation::relu(n, ctx),
@@ -190,8 +240,7 @@ fn dispatch(node: &Node, ctx: &mut ExecutionContext) {
         Node::AveragePool2d(n) => ops::pool::avg_pool2d(n, ctx),
         Node::GlobalAveragePool(n) => ops::pool::global_avg_pool(n, ctx),
         Node::Constant(_) => {}
-        other => {
-            eprintln!("warn: unimplemented op '{}' — skipping", other.name());
-        }
+        // rejected by `is_supported` at load time
+        other => unreachable!("unsupported op '{}'", op_type(other)),
     }
 }

src/interpreter/ops/conv.rs

@@ -1,4 +1,5 @@
 use super::super::context::ExecutionContext;
+use super::padding::{pad2d, resolve_pads, symmetric};
 use crate::tensor::{B, default_device};
 use burn_backend::{
     TensorData, TensorMetadata,
@@ -8,26 +9,30 @@ use burn_backend::{
 use onnx_ir::{
     batch_norm::{BatchNormConfig, BatchNormalizationNode},
     conv2d::Conv2dNode,
-    node::padding::PaddingConfig2d,
 };
 
-fn symmetric_padding(cfg: &PaddingConfig2d) -> [usize; 2] {
-    match cfg {
-        PaddingConfig2d::Valid => [0, 0],
-        PaddingConfig2d::Explicit(top, left, bottom, right) => {
-            [(top + bottom) / 2, (left + right) / 2]
-        }
-    }
-}
-
 pub fn conv2d(node: &Conv2dNode, ctx: &mut ExecutionContext) {
     let x = ctx.resolve(&node.inputs[0]).expect("conv2d: missing input");
     let w = ctx
         .resolve(&node.inputs[1])
         .expect("conv2d: missing weight");
     let bias = node.inputs.get(2).and_then(|a| ctx.resolve(a));
 
-    let pad = symmetric_padding(&node.config.padding);
+    let cfg = &node.config;
+    let pads = resolve_pads(
+        &x,
+        &cfg.padding,
+        &cfg.auto_pad,
+        cfg.kernel_size,
+        cfg.stride,
+        cfg.dilation,
+    );
+    // uneven pads: zero-pad the input up front and run the conv unpadded
+    let (x, pad) = match symmetric(pads) {
+        Some(pad) => (x, pad),
+        None => (pad2d(x, pads, 0.0), [0, 0]),
+    };
+
     let options = ConvOptions::new(
         node.config.stride,
         pad,

src/interpreter/ops/mod.rs

@@ -2,5 +2,6 @@ pub mod activation;
 pub mod conv;
 pub mod elementwise;
 pub mod linear;
+mod padding;
 pub mod pool;
 pub mod reshape;

src/interpreter/ops/padding.rs

@@ -0,0 +1,72 @@
+// Spatial padding shared by conv and pooling ops.
+//
+// Burn's conv/pool kernels only take symmetric padding ([h, w] applied to both sides).
+// ONNX allows asymmetric pads and auto_pad=SAME_*, so we resolve the real per-side
+// pads here and, when they're uneven, pad the input explicitly instead.
+
+use burn_backend::{Slice, TensorData, TensorMetadata, backend::ops::FloatTensorOps};
+use onnx_ir::node::padding::{AutoPad, PaddingConfig2d};
+
+use crate::tensor::{B, FloatPrim, default_device};
+
+/// Per-side pads for an NCHW tensor: [top, left, bottom, right].
+pub type Pads2d = [usize; 4];
+
+/// Resolve the actual pads for a 2D window op, honouring auto_pad.
+pub fn resolve_pads(
+    x: &FloatPrim,
+    padding: &PaddingConfig2d,
+    auto_pad: &AutoPad,
+    kernel: [usize; 2],
+    stride: [usize; 2],
+    dilation: [usize; 2],
+) -> Pads2d {
+    match auto_pad {
+        AutoPad::Valid => [0; 4],
+        AutoPad::SameUpper | AutoPad::SameLower => {
+            let shape = x.shape();
+            let mut begin = [0; 2];
+            let mut end = [0; 2];
+            for i in 0..2 {
+                // output = ceil(input / stride); pad just enough to produce it
+                let input = shape[2 + i];
+                let out = input.div_ceil(stride[i]);
+                let span = (kernel[i] - 1) * dilation[i] + 1;
+                let total = ((out - 1) * stride[i] + span).saturating_sub(input);
+                // SAME_UPPER puts the extra pixel at the end, SAME_LOWER at the start
+                let small = total / 2;
+                (begin[i], end[i]) = match auto_pad {
+                    AutoPad::SameUpper => (small, total - small),
+                    _ => (total - small, small),
+                };
+            }
+            [begin[0], begin[1], end[0], end[1]]
+        }
+        AutoPad::NotSet => match padding {
+            PaddingConfig2d::Valid => [0; 4],
+            PaddingConfig2d::Explicit(top, left, bottom, right) => [*top, *left, *bottom, *right],
+        },
+    }
+}
+
+/// Some([h, w]) if the pads can be passed straight to a Burn kernel.
+pub fn symmetric(pads: Pads2d) -> Option<[usize; 2]> {
+    let [top, left, bottom, right] = pads;
+    (top == bottom && left == right).then_some([top, left])
+}
+
+/// Pad an NCHW tensor's spatial dims with a constant value.
+pub fn pad2d(x: FloatPrim, pads: Pads2d, value: f32) -> FloatPrim {
+    let [top, left, bottom, right] = pads;
+    let shape = x.shape();
+    let (n, c, h, w) = (shape[0], shape[1], shape[2], shape[3]);
+
+    let padded = B::float_from_data(
+        TensorData::full(vec![n, c, h + top + bottom, w + left + right], value),
+        &default_device(),
+    );
+    let range =
+        |start: usize, len: usize| Slice::new(start as isize, Some((start + len) as isize), 1);
+    let slices = [range(0, n), range(0, c), range(top, h), range(left, w)];
+    B::float_slice_assign(padded, &slices, x)
+}

src/interpreter/ops/pool.rs

@@ -1,32 +1,40 @@
 use super::super::context::ExecutionContext;
-use crate::tensor::B;
-use burn_backend::backend::ops::ModuleOps;
+use super::padding::{pad2d, resolve_pads, symmetric};
+use crate::tensor::{B, default_device};
+use burn_backend::{
+    TensorData, TensorMetadata,
+    backend::ops::{FloatTensorOps, ModuleOps},
+};
 use onnx_ir::{
     avg_pool2d::AveragePool2dNode, global_avg_pool::GlobalAveragePoolNode,
-    max_pool2d::MaxPool2dNode, node::padding::PaddingConfig2d,
+    max_pool2d::MaxPool2dNode,
 };
 
-fn symmetric_padding(cfg: &PaddingConfig2d) -> [usize; 2] {
-    match cfg {
-        PaddingConfig2d::Valid => [0, 0],
-        PaddingConfig2d::Explicit(top, left, bottom, right) => {
-            [(top + bottom) / 2, (left + right) / 2]
-        }
-    }
-}
-
 pub fn max_pool2d(node: &MaxPool2dNode, ctx: &mut ExecutionContext) {
     let x = ctx
         .resolve(&node.inputs[0])
         .expect("max_pool2d: missing input");
-    let pad = symmetric_padding(&node.config.padding);
+    let cfg = &node.config;
+    let pads = resolve_pads(
+        &x,
+        &cfg.padding,
+        &cfg.auto_pad,
+        cfg.kernel_size,
+        cfg.strides,
+        cfg.dilation,
+    );
+    // uneven pads: -inf padding never wins the max, so pre-padding is exact
+    let (x, pad) = match symmetric(pads) {
+        Some(pad) => (x, pad),
+        None => (pad2d(x, pads, f32::NEG_INFINITY), [0, 0]),
+    };
     let y = B::max_pool2d(
         x,
-        node.config.kernel_size,
-        node.config.strides,
+        cfg.kernel_size,
+        cfg.strides,
         pad,
-        node.config.dilation,
-        false,
+        cfg.dilation,
+        cfg.ceil_mode,
     );
     ctx.insert(node.outputs[0].name.clone(), y);
 }
@@ -35,15 +43,42 @@ pub fn avg_pool2d(node: &AveragePool2dNode, ctx: &mut ExecutionContext) {
     let x = ctx
         .resolve(&node.inputs[0])
         .expect("avg_pool2d: missing input");
-    let pad = symmetric_padding(&node.config.padding);
-    let y = B::avg_pool2d(
-        x,
-        node.config.kernel_size,
-        node.config.strides,
-        pad,
-        node.config.count_include_pad,
-        false,
+    let cfg = &node.config;
+    let pads = resolve_pads(
+        &x,
+        &cfg.padding,
+        &cfg.auto_pad,
+        cfg.kernel_size,
+        cfg.strides,
+        cfg.dilation,
     );
+    let pool = |t, pad, include_pad| {
+        B::avg_pool2d(
+            t,
+            cfg.kernel_size,
+            cfg.strides,
+            pad,
+            include_pad,
+            cfg.ceil_mode,
+        )
+    };
+
+    let y = match symmetric(pads) {
+        Some(pad) => pool(x, pad, cfg.count_include_pad),
+        None if cfg.count_include_pad => pool(pad2d(x, pads, 0.0), [0, 0], true),
+        None => {
+            // uneven pads that must not count towards the average:
+            // sum(window) / count(real pixels in window), via a padded ones-mask
+            let shape = x.shape();
+            let ones = B::float_from_data(
+                TensorData::full(vec![1, 1, shape[2], shape[3]], 1.0f32),
+                &default_device(),
+            );
+            let sum = pool(pad2d(x, pads, 0.0), [0, 0], true);
+            let count = pool(pad2d(ones, pads, 0.0), [0, 0], true);
+            B::float_div(sum, count)
+        }
+    };
     ctx.insert(node.outputs[0].name.clone(), y);
 }