fix axis normalization

Author: Matthew

src/relay/qnn/op/dequantize.cc

@@ -54,13 +54,13 @@ bool DequantizeRel(const Array<Type>& types, int num_inputs, const Attrs& attrs,
   const auto* dequantize_attrs = attrs.as<DequantizeAttrs>();
   int axis = dequantize_attrs->axis;
   auto rank = static_cast<int>(data->shape.size());
+  axis = (axis < 0) ? ((rank > 0) ? data->shape.size() + axis : 0) : axis;
 
   // If zero point and scale are scalar then axis doesnt matter.
   bool scale_is_scalar = (types[1].as<TensorTypeNode>())->shape.size() == 0;
   bool zp_is_scalar = (types[2].as<TensorTypeNode>())->shape.size() == 0;
 
   if (!(scale_is_scalar && zp_is_scalar)) {
-    axis = (axis < 0) ? ((rank > 0) ? data->shape.size() + axis : 0) : axis;
     ICHECK_LT(axis, rank > 0 ? rank : 1) << "axis " << dequantize_attrs->axis << " is out of range";
     ICHECK_GE(axis, 0) << "axis " << dequantize_attrs->axis << " is out of range";
   }

src/relay/qnn/op/quantize.cc

@@ -52,13 +52,13 @@ bool QuantizeRel(const Array<Type>& types, int num_inputs, const Attrs& attrs,
   const auto* quantize_attrs = attrs.as<QuantizeAttrs>();
   int axis = quantize_attrs->axis;
   auto rank = static_cast<int>(data->shape.size());
+  axis = (axis < 0) ? ((rank > 0) ? data->shape.size() + axis : 0) : axis;
 
   // If zero point and scale are scalar then axis doesnt matter.
   bool scale_is_scalar = (types[1].as<TensorTypeNode>())->shape.size() == 0;
   bool zp_is_scalar = (types[2].as<TensorTypeNode>())->shape.size() == 0;
 
   if (!(scale_is_scalar && zp_is_scalar)) {
-    axis = (axis < 0) ? ((rank > 0) ? data->shape.size() + axis : 0) : axis;
     ICHECK_LT(axis, rank > 0 ? rank : 1) << "axis " << quantize_attrs->axis << " is out of range";
     ICHECK_GE(axis, 0) << "axis " << quantize_attrs->axis << " is out of range";
   }