[GitHub] [tvm] sfvaroglu commented on a diff in pull request #10880: [QNN] Fix qnn.dequantize scale and zp shape

[GitBox <gi...@apache.org>]

sfvaroglu commented on code in PR #10880:
URL: https://github.com/apache/tvm/pull/10880#discussion_r843111159


##########
src/relay/qnn/op/dequantize.cc:
##########
@@ -56,17 +56,20 @@ bool DequantizeRel(const Array<Type>& types, int num_inputs, const Attrs& attrs,
   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 zero point and scale are scalar or have arbitrary rank with one element,
+  // then axis doesn't matter.
+  bool scale_is_scalar = (types[1].as<TensorTypeNode>())->shape.size() == 0 ||
+                         get_const_int((types[1].as<TensorTypeNode>())->Size()) == 1;
+  bool zp_is_scalar = (types[2].as<TensorTypeNode>())->shape.size() == 0 ||
+                      get_const_int((types[2].as<TensorTypeNode>())->Size()) == 1;
 
   if (!(scale_is_scalar && zp_is_scalar)) {
     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";
   }
 
   PrimExpr axis_shape;
-  if (rank > 0) {
+  if (!(scale_is_scalar && zp_is_scalar)) {

Review Comment:
   Done



##########
src/relay/qnn/op/dequantize.cc:
##########
@@ -56,17 +56,20 @@ bool DequantizeRel(const Array<Type>& types, int num_inputs, const Attrs& attrs,
   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 zero point and scale are scalar or have arbitrary rank with one element,
+  // then axis doesn't matter.
+  bool scale_is_scalar = (types[1].as<TensorTypeNode>())->shape.size() == 0 ||
+                         get_const_int((types[1].as<TensorTypeNode>())->Size()) == 1;

Review Comment:
   Yes



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org