Fix: Multiple bug fixes

DeepQuant/CustomForwards/MultiHeadAttention.py

@@ -13,7 +13,7 @@
 
 
 def unrolledQuantMhaForward(
- self: QuantMultiheadAttention, query: Tensor, key: Tensor, value: Tensor
+ self: QuantMultiheadAttention, query: Tensor, key: Tensor, value: Tensor, need_weights: bool = True
 ) -> Tensor:
  """
  Export-friendly forward that explicitly unrolls the multi-head logic.
@@ -84,4 +84,11 @@ def unrolledQuantMhaForward(
  )
 
  attnOutput = self.out_proj(attnOutput)
- return attnOutput
+
+ if need_weights:
+ # return average attention weights over heads
+ attnWeights = attnWeights.view(batchSize, self.num_heads, seqLen, seqLen)
+ attnWeights = attnWeights.sum(dim=1) / self.num_heads
+ else:
+ attnWeights = None
+ return attnOutput, attnWeights

DeepQuant/QuantManipulation/DequantModifier.py

@@ -74,6 +74,12 @@ def unifyLinearDequants(
 
  # Collect and rewire the linear's arguments
  for arg in oldArgs:
+
+ # NEW: skip None and non-Node args, keep them as-is
+ if arg is None or not isinstance(arg, fx.Node):
+ newLinArgs.append(arg)
+ continue
+
  if arg.op == "call_module" and "dequant" in arg.target.lower():
  if "bias_dequant" in arg.target.lower():
  biasDequantNode = arg

DeepQuant/QuantManipulation/ParameterExtractor.py

@@ -45,7 +45,10 @@ def safe_get_scale(quant_obj: Any) -> Any:
  if maybe_scale is None:
  return None
  if isinstance(maybe_scale, torch.Tensor):
- return maybe_scale.item()
+ if maybe_scale.numel() == 1:
+ return maybe_scale.item()
+ else:
+ return maybe_scale
  elif isinstance(maybe_scale, float):
  return maybe_scale
  try:

DeepQuant/QuantManipulation/QuantNodesDivider.py

@@ -132,7 +132,17 @@ def split_quant_nodes(
  for user_node in list(node.users.keys()):
  new_args = []
  for arg in user_node.args:
- new_args.append(dequant_node if arg is node else arg)
+ if arg is node:
+ new_args.append(dequant_node)
+ # if the argument is a tuple or list (e.g. concatenation, addition)
+ elif isinstance(arg, (tuple, list)):
+ seq_type = type(arg)
+ new_seq = []
+ for a in arg:
+ new_seq.append(dequant_node if a is node else a)
+ new_args.append(seq_type(new_seq))
+ else:
+ new_args.append(arg)
  user_node.args = tuple(new_args)
 
  nodes_to_erase.append(node)

Tests/TestConvChannelWise.py

@@ -0,0 +1,59 @@
+# Copyright 2025 ETH Zurich and University of Bologna.
+# Licensed under the Apache License, Version 2.0, see LICENSE for details.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Victor Jung <jungvi@iis.ee.ethz.ch>
+# Federico Brancasi <fbrancasi@ethz.ch>
+
+import pytest
+import torch
+import torch.nn as nn
+import brevitas.nn as qnn
+from brevitas.quant.scaled_int import (
+ Int8ActPerTensorFloat,
+ Int32Bias,
+ Int8WeightPerChannelFloat
+)
+from DeepQuant.ExportBrevitas import exportBrevitas
+
+
+class QuantConvNet(nn.Module):
+
+ convAndLinQuantParams = {
+ "bias": True,
+ "weight_bit_width": 4,
+ "bias_quant": Int32Bias,
+ "input_quant": Int8ActPerTensorFloat,
+ "weight_quant": Int8WeightPerChannelFloat,
+ "output_quant": Int8ActPerTensorFloat,
+ "return_quant_tensor": True,
+ }
+
+ def __init__(self, in_channels: int = 1) -> None:
+ super().__init__()
+ self.inputQuant = qnn.QuantIdentity(return_quant_tensor=True)
+
+ self.conv1 = qnn.QuantConv2d(
+ in_channels=in_channels,
+ out_channels=16,
+ kernel_size=3,
+ padding=1,
+ **QuantConvNet.convAndLinQuantParams
+ )
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+
+ x = self.inputQuant(x)
+ x = self.conv1(x)
+
+ return x
+
+
+@pytest.mark.SingleLayerTests
+def deepQuantTestConv() -> None:
+
+ torch.manual_seed(42)
+
+ model = QuantConvNet().eval()
+ sampleInput = torch.randn(1, 1, 28, 28)
+ exportBrevitas(model, sampleInput, debug=True)

Tests/TestSimpleCNNChannelWise.py

@@ -0,0 +1,108 @@
+# Copyright 2025 ETH Zurich and University of Bologna.
+# Licensed under the Apache License, Version 2.0, see LICENSE for details.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Federico Brancasi <fbrancasi@ethz.ch>
+
+
+import pytest
+import torch
+import torch.nn as nn
+import brevitas.nn as qnn
+from brevitas.quant.scaled_int import (
+ Int8ActPerTensorFloat,
+ Int32Bias,
+ Int8WeightPerChannelFloat,
+)
+from DeepQuant.ExportBrevitas import exportBrevitas
+
+
+class SimpleQuantCNN(nn.Module):
+ """
+ A simple quantized CNN that includes:
+ - Input quantization
+ - Two QuantConv2d layers with Quantized ReLU
+ - MaxPool2d
+ - A final QuantLinear layer
+ """
+
+ convAndLinQuantParams = {
+ "bias": True,
+ "weight_bit_width": 4,
+ "bias_quant": Int32Bias,
+ "input_quant": Int8ActPerTensorFloat,
+ "weight_quant": Int8WeightPerChannelFloat,
+ "output_quant": Int8ActPerTensorFloat,
+ "return_quant_tensor": True,
+ }
+
+ def __init__(self, in_channels: int = 1, num_classes: int = 10) -> None:
+ """
+ Args:
+ in_channels: Number of input channels (e.g., 1 for grayscale).
+ num_classes: Number of output classes for the final linear layer.
+ """
+ super().__init__()
+ self.inputQuant = qnn.QuantIdentity(return_quant_tensor=True)
+
+ self.conv1 = qnn.QuantConv2d(
+ in_channels=in_channels,
+ out_channels=16,
+ kernel_size=3,
+ padding=1,
+ **SimpleQuantCNN.convAndLinQuantParams
+ )
+ self.relu1 = qnn.QuantReLU(bit_width=4, return_quant_tensor=True)
+ self.pool1 = nn.MaxPool2d(kernel_size=2)
+
+ self.conv2 = qnn.QuantConv2d(
+ in_channels=16,
+ out_channels=32,
+ kernel_size=3,
+ padding=1,
+ **SimpleQuantCNN.convAndLinQuantParams
+ )
+ self.relu2 = qnn.QuantReLU(bit_width=4, return_quant_tensor=True)
+ self.pool2 = nn.MaxPool2d(kernel_size=2)
+
+ self.flatten = nn.Flatten()
+ self.fc = qnn.QuantLinear(
+ in_features=32 * 7 * 7, # If input is 28x28, shape after pooling is 7x7
+ out_features=num_classes,
+ **SimpleQuantCNN.convAndLinQuantParams
+ )
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ """
+ Forward pass of the SimpleQuantCNN.
+
+ Args:
+ x: Input tensor of shape [batch_size, in_channels, height, width].
+
+ Returns:
+ A quantized output tensor (batch_size, num_classes).
+ """
+ x = self.inputQuant(x)
+
+ x = self.conv1(x)
+ x = self.relu1(x)
+ x = self.pool1(x)
+
+ x = self.conv2(x)
+ x = self.relu2(x)
+ x = self.pool2(x)
+
+ x = self.flatten(x)
+ x = self.fc(x)
+ return x
+
+
+@pytest.mark.ModelTests
+def deepQuantTestSimpleCNN() -> None:
+
+ torch.manual_seed(42)
+
+ model = SimpleQuantCNN().eval()
+ sampleInput = torch.randn(1, 1, 28, 28)
+
+ exportBrevitas(model, sampleInput, debug=True)