Feat (core): module for runtime computation of exp bias (#1418)

Author: Giuseppe5

src/brevitas/core/bit_width/float.py

@@ -81,10 +81,32 @@ class StaticExponentBias(torch.nn.Module):
         a checkpoint but will be properly handled during device transfers and dtype conversions.
     """
 
-    def __init__(self, exponent_bias, device=None, dtype=None):
+    def __init__(
+            self, exponent_bias: float, device: torch.device = None, dtype: torch.dtype = None):
         super().__init__()
         self.exponent_bias = StatelessBuffer(
             torch.tensor(float(exponent_bias), device=device, dtype=dtype))
 
     def forward(self):
         return self.exponent_bias()
+
+
+class ComputeExponentBias(torch.nn.Module):
+    """
+    Module that returns a runtime-computed exponent bias value.
+
+    Args:
+        exponent_bit_width_impl: Module that returns the exponent bit width
+
+    Examples:
+        >>> exp_bias = ComputeExponentBias(4.)
+        >>> exp_bias()
+        tensor(7.)
+    """
+
+    def __init__(self, exponent_bit_width_impl: torch.nn.Module):
+        super().__init__()
+        self.exponent_bit_width_impl = exponent_bit_width_impl
+
+    def forward(self):
+        return 2 ** (self.exponent_bit_width_impl() - 1) - 1