fixed missing parts in test code causing test failures

Xilinx · Feb 1, 2024 · 382c50b · 382c50b
1 parent 5ad14d4
commit 382c50b
Show file tree

Hide file tree

Showing 5 changed files with 14 additions and 6 deletions.
diff --git a/tests/brevitas/export/quant_module_fixture.py b/tests/brevitas/export/quant_module_fixture.py
@@ -61,7 +61,8 @@
     QuantConv3d,
     QuantConvTranspose1d,
     QuantConvTranspose2d,
-    QuantConvTranspose3d,]
+    QuantConvTranspose3d,
+    ]
 BIT_WIDTHS = [4, 8, 10]  # below 8, equal 8, above 8
 BIAS_BIT_WIDTHS = [8, 16, 32]
 

diff --git a/tests/brevitas/export/test_torch_qcdq.py b/tests/brevitas/export/test_torch_qcdq.py
@@ -36,8 +36,10 @@ def test_torch_qcdq_wbiol_export(
         in_size = (1, IN_CH)
     elif quant_module_impl == QuantConv1d or quant_module_impl == QuantConvTranspose1d:
         in_size = (1, IN_CH, FEATURES)
-    else:
+    elif quant_module_impl == QuantConv2d or quant_module_impl == QuantConvTranspose2d:
         in_size = (1, IN_CH, FEATURES, FEATURES)
+    else:
+        in_size = (1, IN_CH, FEATURES, FEATURES, FEATURES)
 
     inp = torch.randn(in_size)
     quant_module(inp)  # Collect scale factors

diff --git a/tests/brevitas/nn/nn_quantizers_fixture.py b/tests/brevitas/nn/nn_quantizers_fixture.py
@@ -96,7 +96,8 @@
     QuantConv3d,
     QuantConvTranspose1d,
     QuantConvTranspose2d,
-    QuantConvTranspose3d,]
+    QuantConvTranspose3d,
+    ]
 
 ACC_BIT_WIDTHS = [8, 9, 10, 12, 16, 24, 32]
 
@@ -161,8 +162,10 @@ def forward(self, x):
         in_size = (1, IN_CH)
     elif impl in ('QuantConv1d', 'QuantConvTranspose1d'):
         in_size = (1, IN_CH, FEATURES)
-    else:
+    elif impl in ('QuantConv2d', 'QuantConvTranspose2d'):
         in_size = (1, IN_CH, FEATURES, FEATURES)
+    else:
+        in_size = (1, IN_CH, FEATURES, FEATURES, FEATURES)
 
     if input_quantized:
         quant_inp = QuantTensor(

diff --git a/tests/brevitas/nn/test_wbiol.py b/tests/brevitas/nn/test_wbiol.py
@@ -33,7 +33,8 @@
     QuantConv3d,
     QuantConvTranspose1d,
     QuantConvTranspose2d,
-    QuantConvTranspose3d,]
+    QuantConvTranspose3d,
+    ]
 
 
 @pytest_cases.fixture()

diff --git a/tests/brevitas_ort/common.py b/tests/brevitas_ort/common.py
@@ -77,7 +77,8 @@ class A2QWeightQuantizerForTests(Int8AccumulatorAwareWeightQuant):
     QuantConv3d,
     QuantConvTranspose1d,
     QuantConvTranspose2d,
-    QuantConvTranspose3d,]
+    QuantConvTranspose3d,
+    ]
 
 
 def compute_ort(export_name, np_input):