per-tensor-granularity-for-hgq (#44)

* Use config quantization parameter named granularity to determine HGQ granularity, either per-tensor or per-weight

Author: nroope

src/pquant/core/keras/activations.py

@@ -102,6 +102,7 @@ def build(self, input_shape):
                 hgq_gamma=self.hgq_gamma,
                 place="datalane",
                 dynamic_data=self.dynamic_data,
+                granularity=self.config.quantization_parameters.granularity,
             )
         if self.quantize_output:
             self.output_quantizer = Quantizer(
@@ -115,6 +116,7 @@ def build(self, input_shape):
                 hgq_gamma=self.hgq_gamma,
                 place="datalane",
                 dynamic_data=self.dynamic_data,
+                granularity=self.config.quantization_parameters.granularity,
             )
 
         if self.use_multiplier:

src/pquant/core/keras/layers.py

@@ -37,7 +37,15 @@ def __init__(
         self.place = place
         self.granularity = granularity.value if isinstance(granularity, Enum) else granularity
         self.quantizer = create_quantizer(
-            self.k_init, self.i_init, self.f_init, self.overflow, self.round_mode, self.use_hgq, self.is_data, place
+            self.k_init,
+            self.i_init,
+            self.f_init,
+            self.overflow,
+            self.round_mode,
+            self.use_hgq,
+            self.is_data,
+            place,
+            granularity=self.granularity,
         )
         self.is_pretraining = True
         self.hgq_gamma = hgq_gamma
@@ -205,14 +213,41 @@ def get_config(self):
         return config
 
 
-def create_hgq_parameters_quantizer(k, i, f, overflow, round_mode, place, gamma=1e-8):
+def axis_kwargs_for_granularity(granularity, is_data):
+    """Translate a granularity into HGQ's (mutually exclusive) homogeneous/heterogeneous axis spec.
+
+    HGQ only supports per_tensor and per_weight from the granularity enum:
+    - per_tensor: nothing varies -> heterogeneous_axis=() (one bitwidth for the whole tensor)
+    - per_weight: every element varies. For data we keep the batch axis (0) homogeneous via
+      homogeneous_axis=(0,); for weights nothing is shared via homogeneous_axis=().
+
+    per_channel is intentionally NOT supported for HGQ (the channel axis is layout-dependent, so we
+    don't guess it).
+    """
+    if granularity == "per_tensor":
+        return {"heterogeneous_axis": ()}
+    if granularity == "per_weight":
+        return {"homogeneous_axis": (0,) if is_data else ()}
+    if granularity == "per_channel":
+        raise ValueError("per_channel granularity is not supported for HGQ. Use 'per_tensor' or 'per_weight'.")
+    raise ValueError(f"Unsupported granularity: {granularity}")
+
+
+def create_hgq_parameters_quantizer(k, i, f, overflow, round_mode, place, axis_kwargs, gamma=1e-8):
     quantizer_config = QuantizerConfig(
-        q_type="kif", place=place, k0=k, i0=i, f0=f, overflow_mode=overflow, round_mode=round_mode, homogeneous_axis=()
+        q_type="kif",
+        place=place,
+        k0=k,
+        i0=i,
+        f0=f,
+        overflow_mode=overflow,
+        round_mode=round_mode,
+        **axis_kwargs,
     )
     return HGQQuantizer(config=quantizer_config)
 
 
-def create_hgq_data_quantizer(k, i, f, overflow, round_mode, gamma=1e-8):
+def create_hgq_data_quantizer(k, i, f, overflow, round_mode, axis_kwargs, gamma=1e-8):
     quantizer_config = QuantizerConfig(
         q_type="kif",
         place="datalane",
@@ -221,16 +256,19 @@ def create_hgq_data_quantizer(k, i, f, overflow, round_mode, gamma=1e-8):
         f0=f,
         overflow_mode=overflow,
         round_mode=round_mode,
-        homogeneous_axis=(0,),
+        **axis_kwargs,
     )
     return HGQQuantizer(config=quantizer_config)
 
 
-def create_quantizer(k, i, f, overflow, round_mode, is_heterogeneous, is_data, place="datalane", gamma=1e-8):
+def create_quantizer(
+    k, i, f, overflow, round_mode, is_heterogeneous, is_data, place="datalane", granularity="per_weight", gamma=1e-8
+):
     if is_heterogeneous:
+        axis_kwargs = axis_kwargs_for_granularity(granularity, is_data)
         if is_data:
-            return create_hgq_data_quantizer(k, i, f, overflow, round_mode, gamma=gamma)
+            return create_hgq_data_quantizer(k, i, f, overflow, round_mode, axis_kwargs, gamma=gamma)
         else:
-            return create_hgq_parameters_quantizer(k, i, f, overflow, round_mode, place, gamma=gamma)
+            return create_hgq_parameters_quantizer(k, i, f, overflow, round_mode, place, axis_kwargs, gamma=gamma)
     else:
         return get_fixed_quantizer(round_mode=round_mode, overflow_mode=overflow)

src/pquant/core/keras/quantizer.py

@@ -107,6 +107,7 @@ def check_is_built(self, input_shape):
             hgq_gamma=self.hgq_gamma,
             place="datalane",
             dynamic_data=self.dynamic_data,
+            granularity=self.config.quantization_parameters.granularity,
         )
         self.input_quantizer = Quantizer(
             k=self.k_input,
@@ -119,6 +120,7 @@ def check_is_built(self, input_shape):
             hgq_gamma=self.hgq_gamma,
             place="datalane",
             dynamic_data=self.dynamic_data,
+            granularity=self.config.quantization_parameters.granularity,
         )
         if self.use_hgq:
             self.input_quantizer.quantizer.build(input_shape)

src/pquant/core/torch/activations.py

@@ -44,8 +44,12 @@ class HGQQuantizer(nn.Module):
     round_mode : str
         One of 'RND', 'RND_CONV', 'TRN', etc.
     is_data : bool
-        True  → data/activation quantizer (homogeneous over batch axis 0).
-        False → weight/bias quantizer (fully heterogeneous, per-element).
+        True  → data/activation quantizer (batch axis 0 always homogeneous;
+                per_channel keys on channel axis 1).
+        False → weight/bias quantizer (per_channel keys on output-channel axis 0).
+    granularity : str
+        One of 'per_tensor' (one shared bit-width) or 'per_weight' (one per element).
+        Controls which axes are shared. per_channel is not supported.
     gamma : float
         L1 regularisation coefficient on bit-widths.
     i_decay_speed : float
@@ -73,6 +77,7 @@ def __init__(
         overflow_mode: str,
         round_mode: str,
         is_data: bool,
+        granularity: str = "per_weight",
         gamma: float = 1e-8,
         i_decay_speed: float = float("inf"),
         i_min: float = -23.0,
@@ -92,6 +97,7 @@ def __init__(
         self.overflow_mode = overflow_mode.upper()
         self.round_mode = round_mode.upper()
         self.is_data = is_data
+        self.granularity = granularity
         self.gamma = gamma
         self.i_decay_speed = i_decay_speed
         self.i_min = i_min
@@ -129,10 +135,9 @@ def build(self, input_shape: tuple) -> None:
         parameters, not the scalar placeholders from __init__.
         """
         device = self._k.device
+        self.homogeneous_axis = self._homogeneous_axis(len(input_shape))
         bw_shape = self._infer_bw_shape(input_shape)
 
-        self.homogeneous_axis = (0,) if self.is_data else ()
-
         # k: non-trainable sign-bit buffer
         self.register_buffer("_k", torch.full(bw_shape, self.k0, device=device))
 
@@ -146,15 +151,23 @@ def build(self, input_shape: tuple) -> None:
 
         self._built = True
 
+    def _homogeneous_axis(self, ndim: int) -> tuple[int, ...]:
+        """Axes shared (collapsed to 1 in the bit-width tensor), derived from `granularity`.
+
+        HGQ supports only per_tensor and per_weight (data always shares the batch axis 0).
+        per_channel is intentionally unsupported (the channel axis is layout-dependent).
+        """
+        if self.granularity == "per_tensor":
+            return tuple(range(ndim))
+        if self.granularity == "per_weight":
+            return (0,) if self.is_data else ()
+        if self.granularity == "per_channel":
+            raise ValueError("per_channel granularity is not supported for HGQ. Use 'per_tensor' or 'per_weight'.")
+        raise ValueError(f"Unsupported granularity: {self.granularity}")
+
     def _infer_bw_shape(self, input_shape: tuple) -> tuple:
-        """Shape of bit-width parameter tensors given input tensor shape."""
-        if self.is_data:
-            # Batch axis (0) is homogeneous → dimension 0 collapses to 1.
-            shape = list(input_shape)
-            shape[0] = 1
-            return tuple(shape)
-        # Fully heterogeneous (per-parameter): same shape as input.
-        return tuple(input_shape)
+        """Shape of bit-width parameter tensors: homogeneous axes collapse to 1."""
+        return tuple(1 if ax in self.homogeneous_axis else d for ax, d in enumerate(input_shape))
 
     # ------------------------------------------------------------------
     # Properties