Validate TreeEnsemble v5 node references

[Akshay Sonawane (apsonawane)]

This pull request improves the robustness and error handling of the TreeEnsemble implementation in ONNX Runtime. It introduces stricter validation for tree structure attributes, adds cycle detection, and expands the test suite to cover more invalid input scenarios.

Validation and error handling improvements:

• Added checks to ensure tree_roots only contains valid node indices, preventing out-of-range root references.
• Added comprehensive validation for node and leaf indices within the tree traversal, including checks for out-of-range indices and missing required node data.
• Introduced cycle and shared node detection by tracking visited nodes during traversal, enforcing that each internal node is only visited once. [1] [2]

Test coverage enhancements:

• Added the RunInvalidTreeStructureTest helper and new test cases to verify the operator correctly rejects out-of-range roots, out-of-range child nodes, and cycles in the tree structure. [1] [2]

[Copilot]

Pull request overview

This pull request hardens the ai.onnx.ml v5 TreeEnsemble attribute conversion/validation path in ONNX Runtime by adding stricter structural checks (bounds checks + cycle/shared-node detection) and expanding negative tests to ensure malformed tree metadata is rejected deterministically.

Changes:

• Added per-visit validation of node/leaf indices during iterative traversal, plus cycle/shared-internal-node detection via a visited set.
• Added explicit validation that tree_roots references valid indices (preventing out-of-range root access).
• Expanded CPU ML tests with a helper and new failure cases for out-of-range roots/children and cycles.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File	Description
onnxruntime/core/providers/cpu/ml/tree_ensemble_attribute.h	Adds validation and cycle/shared-node detection during TreeEnsemble v5 → v3 conversion, including tree_roots bounds checks.
onnxruntime/test/providers/cpu/ml/tree_ensembler_test.cc	Adds new negative tests (and a helper) that assert invalid tree structures are rejected with expected errors.

Suppressed comments (1)

onnxruntime/core/providers/cpu/ml/tree_ensemble_attribute.h:485

• When the root is detected as a leaf, transformInputOneTree is called with root_id == tree_root (a node index), but the leaf-processing path in transformInputOneTree treats curr_id as an index into leaf_targetids/leaf_weights. For a leaf-only tree in a multi-tree ensemble, the root node index and the referenced leaf index can diverge, which will either throw an out-of-range leaf error or (prior to the new checks) risk reading the wrong leaf weight/target.

Also, the tree_roots bounds check only validates a subset of node attribute arrays; a tree_root can be in-range for nodes_*nodeids but out-of-range for nodes_modes/nodes_featureids, producing a later, less-specific failure.

      ORT_ENFORCE(tree_root >= 0 &&
                      static_cast<uint64_t>(tree_root) < nodes_falsenodeids.size() &&
                      static_cast<uint64_t>(tree_root) < nodes_falseleafs.size() &&
                      static_cast<uint64_t>(tree_root) < nodes_truenodeids.size() &&
                      static_cast<uint64_t>(tree_root) < nodes_trueleafs.size(),
                  "TreeEnsemble tree_roots contains out-of-range node index ", tree_root, ".");
      const size_t tree_root_size_t = static_cast<size_t>(tree_root);
      bool is_leaf = (nodes_falsenodeids[tree_root_size_t] == nodes_truenodeids[tree_root_size_t] &&
                      nodes_falseleafs[tree_root_size_t] && nodes_trueleafs[tree_root_size_t]);
      transformInputOneTree(tree_root_size_t, curr_treeid, 0,
                            is_leaf,
                            membership_values_by_id, output);

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.