Fix issue in per-tensor quantization and missing input_rank (#65)

sony · Oct 5, 2023 · e1d7f11 · e1d7f11
1 parent d708fcf
commit e1d7f11
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 0 deletions.
diff --git a/...rs/pytorch/quantizers/weights_inferable_quantizers/weights_lut_pot_inferable_quantizer.py b/...rs/pytorch/quantizers/weights_inferable_quantizers/weights_lut_pot_inferable_quantizer.py
@@ -155,8 +155,11 @@ def symbolic(g,
             # When None is passed as channel_axis, the op has no attribute of channel_axis,
             # which creates conflict with the onnxruntime function. For this reason, if we quantize
             # per-tensor and channel_axis is None, we set it to 0.
+            # per-tensor and input_rank is None, we set it to 4.
             if not per_channel and channel_axis is None:
                 channel_axis = 0
+            if not per_channel and input_rank is None:
+                input_rank = 4
 
             return g.op(f"{ONNX_CUSTOM_OP_DOMAIN}::WeightsLUTPOTQuantizer", input_tensor,
                         g.op('Constant', value_t=torch.tensor(lut_values, dtype=torch.float32)),

diff --git a/...orch/quantizers/weights_inferable_quantizers/weights_lut_symmetric_inferable_quantizer.py b/...orch/quantizers/weights_inferable_quantizers/weights_lut_symmetric_inferable_quantizer.py
@@ -177,8 +177,11 @@ def symbolic(g,
             # When None is passed as channel_axis, the op has no attribute of channel_axis,
             # which creates conflict with the onnxruntime function. For this reason, if we quantize
             # per-tensor and channel_axis is None, we set it to 0.
+            # per-tensor and input_rank is None, we set it to 4.
             if not per_channel and channel_axis is None:
                 channel_axis = 0
+            if not per_channel and input_rank is None:
+                input_rank = 4
 
             return g.op(f"{ONNX_CUSTOM_OP_DOMAIN}::WeightsLUTSymmetricQuantizer", input_tensor,
                         g.op('Constant', value_t=torch.tensor(lut_values, dtype=torch.float32)),