saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-09 17:28:58 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

check axis range for LayerNorm (#14845)

Browse source 
### Description
<!-- Describe your changes. -->
Add check on axis to make sure it is in a valid range


### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->
This commit is contained in:
Yufeng Li 2023-03-15 14:44:59 -07:00 committed by GitHub
parent 5213546e62
commit da084b0fc1
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 13 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

19
onnxruntime/core/graph/contrib_ops/contrib_defs.cc
View file

@ -24,6 +24,17 @@
#pragma warning(disable : 26426)
#endif
namespace ONNX_NAMESPACE {
inline int64_t HandleNegativeAxis(int64_t axis, int64_t rank) {
if (rank < 0 || axis >= rank || axis < -rank) {
fail_shape_inference("axis ", axis,
" is not in valid range [-", rank, ",", rank - 1, "]");
}
// Handle negative axis
return axis < 0 ? axis + rank : axis;
}
void convPoolShapeInference(
ONNX_NAMESPACE::InferenceContext& ctx,
bool use_dilation, bool require_kernel_shape,
@ -2236,9 +2247,7 @@ void RegisterContribSchemas() {
if (axis_proto) {
axis = axis_proto->i();
}
if (axis < 0) {
axis += input_ndim;
}
axis = HandleNegativeAxis(axis, input_ndim);
if (ctx.getNumOutputs() > 1) {
auto saved_mean_shape = ctx.getOutputType(1)->mutable_tensor_type()->mutable_shape();
@ -2380,9 +2389,7 @@ void RegisterContribSchemas() {
if (axis_proto) {
axis = axis_proto->i();
}
if (axis < 0) {
axis += input_ndim;
}
axis = HandleNegativeAxis(axis, input_ndim);
if (ctx.getNumOutputs() > 1) {
auto saved_inv_std_var_shape = ctx.getOutputType(1)->mutable_tensor_type()->mutable_shape();