qbatchnormdense and permute

Author: calad0i

hls4ml/converters/keras_v3/squark/_base.py

@@ -139,7 +139,7 @@ class SQConvHandler(SQLayerHandler, KV3ConvHandler):
 
 @register
 class SQDenseHandler(SQLayerHandler, KV3DenseHandler):
-    handles = ('squark.layers.core.dense.QDense',)
+    handles = ('squark.layers.core.dense.QDense', 'squark.layers.core.dense.QBatchNormDense')
 
     def handle(
         self,
@@ -148,6 +148,7 @@ def handle(
         out_tensors: Sequence['KerasTensor'],
     ):
         conf = super().handle(layer, in_tensors, out_tensors)
+        conf['class_name'] = 'Dense'
         in_shape: tuple[int, ...] = in_tensors[0].shape[1:]  # type: ignore
         if len(in_shape) > 1:
             if hasattr(layer, 'parallelization_factor'):

hls4ml/model/optimizer/passes/bit_exact.py

@@ -27,6 +27,7 @@
     Pooling2D,
     Reshape,
     Softmax,
+    Transpose,
 )
 from hls4ml.model.optimizer import ModelOptimizerPass, OptimizerPass
 from hls4ml.model.optimizer.passes.hgq_proxy_model import FixedPointQuantizer, UnaryLUT
@@ -154,6 +155,17 @@ def _(layer: Concatenate):
     return ((k0, i0, f0), (k1, i1, f1))
 
 
+@_request_kif.register
+def _(layer: Transpose):
+    k, i, f = requested_kif(layer)
+    perm = layer.attributes['perm']
+    inv_perm = np.argsort(perm)
+    k = np.transpose(k, inv_perm)
+    i = np.transpose(i, inv_perm)
+    f = np.transpose(f, inv_perm)
+    return ((k, i, f),)
+
+
 def requested_kif(layer: Layer) -> KIF_t:
     out_layers = get_output_layers(layer)
     out_shape = get_output_shape(layer)
@@ -319,6 +331,16 @@ def _(layer: Dense):
     return k.astype(np.int8), i, f
 
 
+@_produce_kif.register
+def _(layer: Transpose):
+    k, i, f = get_input_kifs(layer)[0]
+    perm = layer.attributes['perm']
+    k = np.transpose(k, perm)
+    i = np.transpose(i, perm)
+    f = np.transpose(f, perm)
+    return k, i, f
+
+
 def r_im2col(kernel_size: Sequence[int], arr: np.ndarray, buffer: np.ndarray, axis: int):
     w = kernel_size[0]
     if len(kernel_size) == 3:  # 1D