fix axis normalization

fix lint

fix lint again

python/tvm/relay/transform/fake_quantization_to_integer.py

@@ -139,9 +139,8 @@ def conv2d(expr, type_map):
     out = relay.qnn.op.conv2d(
         x, weight, x_t.zero_point, w_t.zero_point, x_t.scale, w_t.scale, **attrs
     )
-    scale_shape = infer_shape(conv_scale)
     out_layout = attrs["out_layout"] if attrs["out_layout"] != "" else attrs["data_layout"]
-    out_axis = tvm.tir.bijective_layout(out_layout, "NCHW").backward_index(list(range(4)))[1]
+    out_axis = bijective_layout(out_layout, "NCHW").backward_index(list(range(4)))[1]
     return [out, TensorAffineType(conv_scale, conv_zp, out.attrs.out_dtype, out_axis.value)]
 
 
@@ -252,6 +251,7 @@ def clip(expr, type_map):
 
 @register_fake_quantization_to_integer("nn.relu")
 def relu(expr, type_map):
+    """Rewrite a relu op"""
     arg = expr.args[0]
     t = type_map[arg]
     scale_shape = infer_shape(t.scale)

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";
   }