Fix (export/qonnx): shape propagation in custom ops

Xilinx · Nov 26, 2023 · 3df5396 · 3df5396
1 parent 740d9d4
commit 3df5396
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/src/brevitas/export/onnx/qonnx/function.py b/src/brevitas/export/onnx/qonnx/function.py
@@ -3,6 +3,7 @@
 
 import torch
 from torch.autograd import Function
+from torch.onnx.symbolic_helper import _get_tensor_sizes
 
 from brevitas.core.bit_width import BitWidthConst
 from brevitas.core.function_wrapper.clamp import TensorClamp
@@ -19,6 +20,7 @@ class BrevitasBinaryQuantFn(Function):
     @staticmethod
     def symbolic(g, x, scale, zero_point, bit_width, narrow_range, signed, rounding_mode):
         ret = g.op(f'{DOMAIN_STRING}::BipolarQuant', x, scale)
+        ret.setType(x.type())
         return ret
 
     @staticmethod
@@ -40,6 +42,7 @@ def symbolic(g, x, scale, zero_point, bit_width, narrow_range, signed, rounding_
             rounding_mode_s=rounding_mode,
             signed_i=int(signed),
             narrow_i=int(narrow_range))
+        ret.setType(x.type())
         return ret
 
     @staticmethod
@@ -66,6 +69,7 @@ def symbolic(g, x, scale, zero_point, input_bit_width, output_bit_width, roundin
             input_bit_width,
             output_bit_width,
             rounding_mode_s=rounding_mode)
+        ret.setType(x.type())
         return ret
 
     @staticmethod