Fix oob dereference

Merged
AP
5mo ago
+28-0

This pull request strengthens shape inference validation for several custom BERT-related ONNX operators by adding explicit rank checks for input tensors. These changes ensure that input tensors meet minimum rank requirements, improving error messaging and preventing incorrect shape propagation.

Enhanced shape validation for custom ONNX operators:

RelativePositionBias and GatedRelativePositionBias:

  • Added checks to ensure bias_table (for RelativePositionBias) and token_offset (for GatedRelativePositionBias) inputs have rank ≥ 2, with clear error messages if not. [1] [2]

CausalConvWithState:

  • Added checks to ensure both input and weight tensors have rank ≥ 2, failing shape inference with descriptive errors if violated.

LinearAttention:

  • Added checks to ensure query and value tensors have rank ≥ 3 for both output and state shape inference, with early returns or errors if requirements are not met. [1] [2]

SkipLayerNormalization:

  • Added a check to ensure the input tensor has rank ≥ 1, improving error reporting for invalid input shapes.