Fix zero point

Xilinx · Dec 18, 2024 · e7ef917 · e7ef917
1 parent d7b5036
commit e7ef917
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git a/src/brevitas/quant_tensor/groupwise_float_quant_tensor.py b/src/brevitas/quant_tensor/groupwise_float_quant_tensor.py
@@ -116,8 +116,9 @@ def expand(self):
                 torch.unbind(new_value, dim=self.group_dim)[residual:], dim=self.group_dim)
             new_scale = torch.stack(
                 torch.unbind(new_scale, dim=self.group_dim)[residual:], dim=self.group_dim)
-            new_zp = torch.stack(
-                torch.unbind(new_zp, dim=self.group_dim)[residual:], dim=self.group_dim)
+            if self.zero_point_.shape != ():
+                new_zp = torch.stack(
+                    torch.unbind(new_zp, dim=self.group_dim)[residual:], dim=self.group_dim)
 
         return new_value, new_scale, new_zp
 

diff --git a/src/brevitas/quant_tensor/groupwise_int_quant_tensor.py b/src/brevitas/quant_tensor/groupwise_int_quant_tensor.py
@@ -102,8 +102,9 @@ def expand(self):
                 torch.unbind(new_value, dim=self.group_dim)[residual:], dim=self.group_dim)
             new_scale = torch.stack(
                 torch.unbind(new_scale, dim=self.group_dim)[residual:], dim=self.group_dim)
-            new_zp = torch.stack(
-                torch.unbind(new_zp, dim=self.group_dim)[residual:], dim=self.group_dim)
+            if self.zero_point_.shape != ():
+                new_zp = torch.stack(
+                    torch.unbind(new_zp, dim=self.group_dim)[residual:], dim=self.group_dim)
 
         return new_value, new_scale, new_zp