fix exact bits format

src/HGQ/layers/base.py

@@ -34,10 +34,7 @@ def input_bw(self):
     @property
     def input_bw_exact(self):
         assert len(self._inbound_nodes) <= 1, f"Layer {self.name} is reused {len(self._inbound_nodes)} times. This is not allowed."
-        try:
-            return self.last_layer.act_bw_exact.astype(np.float32)
-        except AssertionError:
-            return None
+        return self.last_layer.act_bw_exact.astype(np.float32)
 
 
 class HLayerBase(ABSBaseLayer):
@@ -195,10 +192,11 @@ def act_bw(self):
         return tf.broadcast_to(bw, (1,) + self.output_shape[1:])
 
     @property
-    def act_bw_exact(self):
+    def act_bw_exact(self) -> np.ndarray:
         """Returns the exact bitwidth of the pre-activation values. Non-differentiable. For post-training use."""
         kn, int_bits, fb = self.paq.get_bits_exact(pos_only=self._relu_act)
-        return int_bits + fb + kn
+        bw = int_bits + fb + kn
+        return np.broadcast_to(bw, (1,) + self.output_shape[1:])
 
     @property
     def fused_bias(self):

src/HGQ/layers/passive_layers.py

@@ -23,6 +23,12 @@ def act_bw(self):
         input_bw = tf.reshape(self.input_bw, shape)
         return tf.ensure_shape(input_bw, shape)
 
+    @property
+    def act_bw_exact(self) -> np.ndarray:
+        """Returns the bitwidth of the pre-activation values. Differentiable."""
+        shape = (1,) + self.output_shape[1:]
+        return self.input_bw_exact.reshape(shape)
+
 
 @register_keras_serializable(package="HGQ")
 class Signature(PLayerBase):
@@ -55,6 +61,10 @@ def call(self, x, training=None, record_minmax=None):
     def act_bw(self):
         return tf.keras.backend.cast_to_floatx(self.bits)
 
+    @property
+    def act_bw_exact(self) -> np.ndarray:
+        return self.bits.numpy()  # type: ignore
+
     @property
     def input_bw(self):
         raise ValueError('Signature layer does not have input_bw')