[Akshay Sonawane (apsonawane)]

This pull request improves input validation for category set IDs in the TreeEnsemble operators and adds corresponding unit tests to ensure invalid category set IDs are properly rejected. The key changes include stricter runtime validation and new tests for both classifier and regressor variants.

Input validation improvements:

• Added a new method ValidateCategorySetId in tree_ensemble_common.h to ensure category set IDs are finite, non-negative integers within bounds, and updated GetCategorySet to use this validation.
• During model initialization, category set IDs for nodes using BRANCH_MEMBER_BIGSET are now validated for correctness.

Testing enhancements:

• Added a unit test TreeEnsembleClassifierRejectsInvalidCategorySetId to verify that the classifier operator fails when given an invalid category set ID.
• Added a unit test TreeEnsembleRegressorRejectsInvalidCategorySetId to verify that the regressor operator fails when given an invalid category set ID.

Other:

• Included <cmath> header to support new validation logic.

[Copilot]

🟡 Changes recommended

The new tests do not detect removal of integer-ID validation, leaving fractional-ID truncation and other invalid category-set IDs untested with a nonempty collection.

1 open issue

• Add ID validation tests with nonempty category sets · New

What changed in this PR

This PR strengthens category-set index validation in ONNX Runtime’s shared CPU TreeEnsemble implementation, rejecting invalid IDs before category-set access.

Changes:

• Require category-set IDs to be finite, non-negative integers within bounds.
• Validate BRANCH_MEMBER_BIGSET IDs during initialization and category-set lookup.
• Add classifier and regressor regression tests for ID 0 with an empty category-set collection.

File	Description
onnxruntime/​test/​providers/​cpu/​ml/​treeregressor_test.cc	Adds a float-threshold regressor rejection test.
onnxruntime/​test/​providers/​cpu/​ml/​tree_ensembler_classifier_test.cc	Adds a double-threshold classifier rejection test.
onnxruntime/​core/​providers/​cpu/​ml/​tree_ensemble_common.h	Adds shared ID validation during initialization and lookup.

_{Review effort: Balanced}

[Ti-Tai Wang (titaiwangms)]

No blocking issues found in a bounded static review of the diff, surrounding initialization/category-set access, and added tests. The validation precedes category-set indexing. One non-blocking test-coverage suggestion is inline.

Validation scope: static inspection only; I did not build or run the tests or sanitizers.

[Copilot]

🟡 Changes recommended

The float-based bounds check can reject a valid final category-set ID for large counts, preventing a valid opset-5 TreeEnsemble from initializing.

3 open issues

• Float threshold count rounding rejects the last valid category-set ID · New
• Avoid repeating category-set ID validation in the inference hot loop · New
• Add ID validation tests with nonempty category sets

_{Review effort: Max}

[Copilot]

🟡 Changes recommended

The new rejection tests abort under ORT_NO_EXCEPTIONS, and both operator tests assert error text that does not match the emitted diagnostic.

2 open issues

• Align operator-test assertions with the bounds-check message · New
• Guard rejection tests when exceptions are disabled · New

3 resolved since last review

• Avoid repeating category-set ID validation in the inference hot loop
• Float threshold count rounding rejects the last valid category-set ID
• Add ID validation tests with nonempty category sets

_{Review effort: Max}

[Tianlei Wu (tianleiwu)]

Reviewed current head 47a803c4cbdb31b11a595c47b182bf6ae22014d6. No blocking code findings. The ID checks run before node construction, guard the floating-point-to-size_t conversion, and compare against the original integer set count. The inference lookup remains unchanged. The earlier feedback on float-count rounding, hot-loop overhead, diagnostics, and exception-disabled tests is addressed.

One nonblocking description correction: the PR body still says GetCategorySet was updated to use validation. In the final implementation, validation happens during Init and GetCategorySet retains its direct lookup. Please update that bullet to match the code.

Validation performed:

• Compiled all three modified test translation units with GCC, both normally and with ORT_NO_EXCEPTIONS, ONNX_NO_EXCEPTIONS, and -fno-exceptions.
• Built a focused C++ harness against this head's header, linking existing local ORT CPU dependency archives: 78 checks passed. These exercised float/double IDs, invalid/nonfinite/fractional values, empty and nonempty sets, the size_t boundary, the 16,777,217-set float-count boundary, and initialization plus membership traversal using both threshold storage paths.
• This was focused validation; I did not run a clean full build or the full provider-test suite.

The remaining CI failures appear outside this diff: Web CI reports detached ArrayBuffer errors in browser tests; WebGPU plugin tests report DXGI_ERROR_DEVICE_REMOVED; GPU documentation validation reports generated-document differences.