support more general case

Author: calad0i

hls4ml/converters/keras_v3/merge.py

@@ -39,7 +39,8 @@ def handle(
         match cls_name:
             case 'Concatenate':
                 rank = len(output_shape)
-                class_name = f'Concatenate{rank}d'
+                class_name = 'Concatenate'
+                op = f'Concatenate{rank}d'
                 config['axis'] = layer.axis
             case 'Dot':
                 msg = (

hls4ml/model/optimizer/passes/bit_exact.py

@@ -3,7 +3,7 @@
 
 import re
 import typing
-from collections.abc import Generator, Sequence
+from collections.abc import Sequence
 from copy import copy
 from functools import reduce, singledispatch
 from math import ceil, log2, prod
@@ -220,8 +220,16 @@ def _produce_kif(layer: Layer) -> KIF_t:
 
 @_produce_kif.register
 def _(layer: Input):
-    k = np.ones(get_output_shape(layer), dtype=np.int8)
-    i = f = np.full(get_output_shape(layer), 126, dtype=np.int8)
+    shape = get_output_shape(layer)
+    if layer.attributes.get('trusted', False):
+        precision: FixedPrecisionType = layer.get_output_variable().type.precision
+        k, i, f = precision.signed, precision.integer - precision.signed, precision.fractional
+        k = np.full(shape, k, dtype=np.int8)
+        i = np.full(shape, i, dtype=np.int8)
+        f = np.full(shape, f, dtype=np.int8)
+    else:
+        k = np.ones(shape, dtype=np.int8)
+        i = f = np.full(shape, 126, dtype=np.int8)
     return k, i, f
 
 
@@ -603,8 +611,8 @@ def kif_arrs_to_ints(arr: tuple[np.ndarray, np.ndarray, np.ndarray]):
     return tuple(int(np.max(a)) for a in arr)
 
 
-def produce_kif(layer: Layer) -> KIF_t:
-    if layer.attributes.get('_produce_kif'):
+def produce_kif(layer: Layer, force_reset=False) -> KIF_t:
+    if layer.attributes.get('_produce_kif') and not force_reset:
         return layer.attributes['_produce_kif']
     kif = _produce_kif(layer)
     layer.attributes['_produce_kif'] = kif
@@ -849,7 +857,9 @@ def transform(self, model: 'ModelGraph'):
         for node in model.graph.values():
             if node.attributes.get('bit_exact_transformed'):
                 continue
-            produce_kif(node)  # Shrink FixedPointQuantizer bits when possible to be used in backward flow (requested_kif).
+            produce_kif(
+                node, force_reset=True
+            )  # Shrink FixedPointQuantizer bits when possible to be used in backward flow (requested_kif).
 
         for node in model.graph.values():
             if node.attributes.get('bit_exact_transformed'):
@@ -858,22 +868,29 @@ def transform(self, model: 'ModelGraph'):
             node.attributes['bit_exact_transformed'] = True
 
         for node in model.graph.values():
-            if node.attributes.get('_produce_kif'):
+            if '_produce_kif' in node.attributes:
                 del node.attributes['_produce_kif']
-            if node.attributes.get('_request_kif'):
+            if '_request_kif' in node.attributes:
                 del node.attributes['_request_kif']
 
         return True
 
 
-def get_output_quantizers(node: Layer) -> Generator[FixedPointQuantizer, None, None]:
+def get_output_layers_and_quantizers(
+    node: Layer, layers: list | None = None, quantizers: list | None = None
+) -> tuple[list[Layer], list[FixedPointQuantizer]]:
+
+    layers = layers if layers is not None else []
+    quantizers = quantizers if quantizers is not None else []
     for _node in get_output_layers(node):
         if isinstance(_node, FixedPointQuantizer):
-            yield _node
-        elif isinstance(_node, (Reshape, Transpose)):
-            yield from get_output_quantizers(_node)
+            quantizers.append(_node)
+        elif isinstance(_node, (Reshape, Transpose, Concatenate)):
+            layers.append(_node)
+            get_output_layers_and_quantizers(_node, layers, quantizers)
         else:
             raise ValueError(f'Layer {node.name} ({node.class_name}) unexpected input layer chain.')
+    return layers, quantizers
 
 
 class FixInputPrecision(OptimizerPass):
@@ -885,17 +902,17 @@ def match(self, node: Layer):
         return node.get_output_variable().type.precision.width > 100
 
     def transform(self, model, node: Layer):
-        out_layers = list(get_output_quantizers(node))
+        layers, out_quantizers = get_output_layers_and_quantizers(node)
 
-        if len(out_layers) == 0:  # Input connected to nothing
+        if len(out_quantizers) == 0:  # Input connected to nothing
             new_type = to_hls4ml_fixed(0, 0, 1, f'{node.name}_t')
             node.get_output_variable().type = new_type
             node.model.config.layer_name_precision[node.name] = str(new_type)
             return False
 
-        sat_modes = [l.SAT for l in out_layers]
+        sat_modes = [l.SAT for l in out_quantizers]
         sat_modes_set = set(sat_modes)
-        rnd_modes = [l.RND for l in out_layers]
+        rnd_modes = [l.RND for l in out_quantizers]
         rnd_modes_set = set(rnd_modes)
         illegal_sat_modes = sat_modes_set - {'WRAP', 'SAT', 'SAT_SYM'}
         illegal_rnd_modes = rnd_modes_set - {'TRN', 'RND'}
@@ -906,7 +923,7 @@ def transform(self, model, node: Layer):
         if illegal_rnd_modes:
             warn(f'Saturation mode {illegal_rnd_modes} may compromise bit-exactness. Forcing at maximum 24 fractional bits.')
 
-        kifs = [_produce_kif(l) for l in out_layers]
+        kifs = [_produce_kif(l) for l in out_quantizers]
         i = np.max([np.max(i) for _, i, _ in kifs])
         k = np.max([np.max(k) for k, _, _ in kifs])
         if illegal_rnd_modes:
@@ -921,4 +938,15 @@ def transform(self, model, node: Layer):
             new_type.precision.saturation_mode = 'SAT'
         node.get_output_variable().type = new_type
         node.model.config.layer_name_precision[node.name] = str(new_type)
+        node.attributes['trusted'] = True
+
+        for layer in layers:
+            produce_kif(layer, force_reset=True)
+        for layer in layers:
+            register_precision(layer)
+        for layer in layers:
+            if '_produce_kif' in layer.attributes:
+                del layer.attributes['_produce_kif']
+            if '_request_kif' in layer.attributes:
+                del layer.attributes['_request_kif']
         return False

hls4ml/model/optimizer/passes/hgq_proxy_model.py

@@ -6,7 +6,7 @@
 import numpy as np
 
 from hls4ml.model.attributes import Attribute, TypeAttribute, WeightAttribute
-from hls4ml.model.layers import Activation, Layer, Reshape, register_layer
+from hls4ml.model.layers import Activation, Layer, Reshape, Transpose, register_layer
 from hls4ml.model.optimizer import OptimizerPass, register_pass
 from hls4ml.model.types import FixedPrecisionType, UnspecifiedPrecisionType
 
@@ -97,7 +97,7 @@ def propagate(self, node: Layer, precision: FixedPrecisionType):
         node.attributes['result_t'].precision = precision
         node.attributes['_result_t_propagated'] = True
 
-        if not isinstance(node, Reshape):
+        if not isinstance(node, (Reshape, Transpose)):
             return node
 
         inp_layer = get_input_layers(node)[0]