fix(reference): accept scalar and omitted optional inputs in NonMaxSuppression - #8655

Description

NonMaxSuppression has three optional inputs, max_output_boxes_per_class, iou_threshold and score_threshold. The spec calls each of them a scalar, and gives max_output_boxes_per_class and iou_threshold a default of 0. The reference implementation in onnx/reference/ops/op_non_max_suppression.py fails in both cases:

• It reads the values with [0] (lines 137, 140 and 143 before this change), so a 0-D tensor raises IndexError. Only 1-D tensors holding one element work, and those are what all the existing node tests use.
• It calls .size on max_output_boxes_per_class and iou_threshold without checking for None (lines 159 and 161), so leaving out either input raises AttributeError. score_threshold already had that check.
• The early return meant for max_output_boxes_per_class == 0 tests .size on the clamped value (line 201). That value is a Python int when the input is empty or negative, so the check raises AttributeError there and never returns otherwise.

On current main:

import numpy as np
from onnx.helper import make_node
from onnx.reference import ReferenceEvaluator

boxes = np.array(
    [[[0.0, 0.0, 1.0, 1.0], [0.0, 0.1, 1.0, 1.1], [0.0, -0.1, 1.0, 0.9],
      [0.0, 10.0, 1.0, 11.0], [0.0, 10.1, 1.0, 11.1], [0.0, 100.0, 1.0, 101.0]]],
    dtype=np.float32,
)
scores = np.array([[[0.9, 0.75, 0.6, 0.95, 0.5, 0.3]]], dtype=np.float32)

# 1. Scalar (0-D) inputs.
# IndexError: too many indices for array: array is 0-dimensional, but 1 were indexed
node = make_node("NonMaxSuppression", ["boxes", "scores", "max", "iou", "score"], ["selected"])
ReferenceEvaluator(node).run(None, {
    "boxes": boxes, "scores": scores, "max": np.array(3, dtype=np.int64),
    "iou": np.array(0.5, dtype=np.float32), "score": np.array(0.4, dtype=np.float32),
})

# 2. iou_threshold and score_threshold omitted.
# TypeError, caused by AttributeError: 'NoneType' object has no attribute 'size'
node = make_node("NonMaxSuppression", ["boxes", "scores", "max"], ["selected"])
ReferenceEvaluator(node).run(None, {"boxes": boxes, "scores": scores, "max": np.array([3], dtype=np.int64)})

With this change the two calls return [[0, 0, 3], [0, 0, 0]] and [[0, 0, 3], [0, 0, 0], [0, 0, 5]]. With no optional inputs at all, the result is an empty [0, 3] tensor.

The fix adds a small helper that turns each present, non-empty optional input into a 1-D array and treats a missing or empty one as absent, so both 0-D and single-element 1-D tensors work and the defaults apply. The early return now compares the value with 0 and returns a [0, 3] tensor, the documented output shape. Nothing changes for inputs that already worked.

Motivation and Context

A model that passes these inputs as 0-D tensors, which is how the spec describes them, cannot run in ReferenceEvaluator. The same is true of a model that leaves out iou_threshold and relies on its default. The node tests did not catch this because every one of them passes all five inputs as 1-D tensors.

Testing

Added to tests/python/reference_evaluator_test.py:

• test_non_max_suppression_scalar_inputs, run with 0-D and with 1-D inputs. The 1-D case is a control that already passed.
• test_non_max_suppression_omitted_optional_inputs, which covers all three optional inputs omitted, both thresholds omitted, and only iou_threshold omitted.

Added the node test test_nonmaxsuppression_scalar_inputs. It uses the same data as test_nonmaxsuppression_suppress_by_IOU_and_scores, with the three inputs as 0-D tensors. docs/Operators.md and docs/TestCoverage.md were regenerated with python onnx/defs/gen_doc.py and python onnx/backend/test/stat_coverage.py (ONNX_ML=1). The only change in either file is the new example.

$ python -m pytest tests/python/reference_evaluator_test.py -k non_max_suppression -q
5 passed, 523 deselected

$ python -m pytest tests/python/reference_evaluator_test.py -q
524 passed, 4 skipped

$ python -m pytest tests/python/backend_reference_test.py -q -k nonmaxsuppression
11 passed, 11 skipped, 4092 deselected

$ python -m pytest tests/python/backend_reference_test.py -q
2030 passed, 2084 skipped

On unmodified main, the four new reference evaluator cases fail: the 0-D case with IndexError, the omitted-input cases with AttributeError. The 1-D control passes on both. The new node test fails on main with the same IndexError. ruff check and ruff format --check pass on the changed files.

I also ran the new node test model through onnxruntime 1.31.0 as an independent implementation, and it returns the expected [[0, 0, 3], [0, 0, 0]]. In 800 randomized cases (400 with 0-D inputs and 400 with 1-D inputs, covering both box formats and tied scores) the patched reference matches onnxruntime exactly.

🤖 Generated with Claude Code