Fix out of bounds read issue in cross.entropy.cc

[vraspar]

Description

Add bounds checking for label tensor values in SparseSoftmaxCrossEntropy::Compute to prevent out-of-bounds memory reads.

The SparseSoftmaxCrossEntropy operator uses label_data[i] (int64_t) directly as an array index into the log-probability buffer without validating that the value falls within [0, D) where D is the number of classes. A malicious ONNX model can embed arbitrary label values in a model initializer, causing the operator to read heap memory beyond the log-probability buffer.

Affected expressions in cross_entropy.cc:

loss_sample[i] = -log_prob_data[i * d + label_data[i]] * weight_data[i];  // weighted path
loss_sample[i] = -log_prob_data[i * d + label_data[i]];                   // unweighted path

Existing shape validation confirms label and logit dimensions are compatible, but never validates label values against the class dimension.

Fix

Added a validation loop before the loss computation that returns an error status if any label value is outside [0, D):

for (ptrdiff_t i = 0; i < n; i++) {
  ORT_RETURN_IF(label_data[i] < 0 || label_data[i] >= d,
                "SparseSoftmaxCrossEntropy: label value ", label_data[i],
                " at index ", i, " is out of range [0, ", d, ")");
}

[Copilot]

Pull request overview

This PR adds bounds checking for label tensor values in SparseSoftmaxCrossEntropy::Compute to prevent out-of-bounds memory reads. Without this fix, a malicious ONNX model could embed arbitrary int64 label values that would be used directly as array indices into the log-probability buffer (which has size N×D), allowing heap memory reads beyond that buffer.

Changes:

• A validation loop over all n label values is added after the softmax computation to return an error if any label falls outside [0, D)
• The check uses ORT_RETURN_IF with a descriptive message that includes the offending label value and its index

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.