Update dense_rewrite

apache · Apr 9, 2019 · 89f1154 · 89f1154
1 parent 6a6082b
commit 89f1154
Showing 1 changed file with 7 additions and 5 deletions.
diff --git a/python/tvm/relay/quantize/_annotate.py b/python/tvm/relay/quantize/_annotate.py
@@ -167,12 +167,14 @@ def dense_rewrite(ref_call, new_args, ctx):
     lhs_expr, lhs_kind = _get_expr_kind(new_args[0])
     rhs_expr, rhs_kind = _get_expr_kind(new_args[1])
 
-    if lhs_kind != QAnnotateKind.INPUT and rhs_kind is None:
+    if lhs_kind is None or lhs_kind != QAnnotateKind.INPUT:
         lhs_expr = attach_simulated_quantize(lhs_expr, QAnnotateKind.INPUT)
-        rhs_expr = attach_simulated_quantize(rhs_expr, QAnnotateKind.WEIGHT)
-        expr = _forward_op(ref_call, [lhs_expr, rhs_expr])
-        return QAnnotateExpr(expr, QAnnotateKind.ACTIVATION)
-    return None
+
+    assert rhs_kind is None
+    rhs_expr = attach_simulated_quantize(rhs_expr, QAnnotateKind.WEIGHT)
+
+    expr = _forward_op(ref_call, [lhs_expr, rhs_expr])
+    return QAnnotateExpr(expr, QAnnotateKind.ACTIVATION)
 
 
 @register_annotate_function("multiply")