From da084b0fc163f131c117ceffe654adef992d6c63 Mon Sep 17 00:00:00 2001
From: Yufeng Li <liyufeng1987@gmail.com>
Date: Wed, 15 Mar 2023 14:44:59 -0700
Subject: [PATCH] check axis range for LayerNorm (#14845)

### 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. -->
---
 .../core/graph/contrib_ops/contrib_defs.cc    | 19 +++++++++++++------
 1 file changed, 13 insertions(+), 6 deletions(-)

diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc
index e80841ef63..5ad939c8b8 100644
--- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc
+++ b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc
@@ -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();