sequential and i/o tensor name parsing fix

Author: calad0i

hls4ml/converters/keras_v3/_base.py

@@ -1,5 +1,19 @@
 import typing
-from typing import Any, Callable, Sequence
+from types import FunctionType
+from typing import Any, Callable, Sequence, TypedDict
+
+
+class DefaultConfig(TypedDict, total=False):
+    name: str
+    class_name: str
+    module: str
+    input_keras_tensor_names: list[str]
+    input_shape: list[list[int]]
+    output_keras_tensor_names: list[str]
+    epsilon: float
+    use_bias: bool
+    data_format: str
+
 
 if typing.TYPE_CHECKING:
     import keras
@@ -49,7 +63,7 @@ def deco(func: T_kv3_handler):
     return deco
 
 
-def maybe_add_attrs(config: dict[str, Any], obj: Any, *attrs: str):
+def maybe_add_attrs(config: dict[str, Any] | DefaultConfig, obj: Any, *attrs: str):
     for attr in attrs:
         if attr not in config and hasattr(obj, attr):
             config[attr] = getattr(obj, attr)
@@ -103,36 +117,55 @@ def __call__(
         """  # noqa: E501
         import keras
 
-        config0 = self.handle(layer, in_tensors, out_tensors)
-        if isinstance(config0, tuple):
-            return config0
-
         name = layer.name
         class_name = layer.__class__.__name__
         module = layer.__module__
-        config1 = {
+
+        default_config: DefaultConfig = {
             'name': name,
             'class_name': class_name,
             'module': module,
             'input_keras_tensor_names': [t.name for t in in_tensors],
-            'input_shape': [list(t.shape[1:]) for t in in_tensors],
+            'input_shape': [list(t.shape[1:]) for t in in_tensors],  # type: ignore
             'output_keras_tensor_names': [t.name for t in out_tensors],
         }
 
-        maybe_add_attrs(config1, layer, 'epsilon', 'use_bias', 'data_format')
+        maybe_add_attrs(default_config, layer, 'epsilon', 'use_bias', 'data_format')
 
-        config1.update(config0)
-        ret = (config1,)
+        mandatory_keys = ['name', 'class_name', 'output_keras_tensor_names', 'input_keras_tensor_names']
 
+        self.default_config = default_config
+        config0 = self.handle(layer, in_tensors, out_tensors)
+        del self.default_config
+
+        if isinstance(config0, tuple):
+            for conf in config0:
+                for key in mandatory_keys:
+                    assert key in conf, f"Key {key} missing from layer {name} handled by {self.__class__.__name__}"
+            return config0
+
+        config = {}
+        config.update(default_config)
+        config.update(config0)
+        ret = (config,)
+
+        # If activation exists, append it
         activation = getattr(layer, 'activation', None)
         if activation not in (keras.activations.linear, None):
-            act_cls_name = activation.__class__.__name__
+            assert len(out_tensors) == 1, f"Layer {name} has more than one output, but has an activation function"
+            assert isinstance(activation, FunctionType), f"Activation function for layer {name} is not a function"
+            intermediate_tensor_name = f'{out_tensors[0].name}_activation'
+            ret[0]['output_keras_tensor_names'] = [intermediate_tensor_name]
+            act_cls_name = activation.__name__
             act_config = {
                 'class_name': 'Activation',
                 'activation': act_cls_name,
                 'name': f'{name}_{act_cls_name}',
+                'input_keras_tensor_names': [intermediate_tensor_name],
+                'output_keras_tensor_names': [out_tensors[0].name],
             }
             ret = *ret, act_config
+
         return ret
 
     def handle(

hls4ml/converters/keras_v3_to_hls.py

@@ -1,5 +1,6 @@
 import typing
 from itertools import chain
+from types import FunctionType
 from typing import Any, Callable, Sequence
 
 if typing.TYPE_CHECKING:
@@ -154,7 +155,10 @@ def v2_call(
         self, layer: 'keras.layers.Layer', inp_tensors: Sequence['KerasTensor'], out_tensors: Sequence['KerasTensor']
     ):
         # keras v2 handlers fallback
-        print("v2 handler")
+        print(f"v2 handler used for layer {layer.name}")
+
+        import keras
+
         config = layer.get_config()
         layer_dict = {'config': config, 'class_name': layer.__class__.__name__}
 
@@ -176,16 +180,22 @@ def get_weights_data(self, layer_name, var_name):
             return None
 
         ret, _ = handler(layer_dict, input_names, input_shapes, reader)
-        ret['outputs'] = output_names
+        ret['output_keras_tensor_names'] = output_names
+        ret['input_keras_tensor_names'] = input_names
         ret = (ret,)
 
         activation = getattr(layer, 'activation', None)
         if activation not in (keras.activations.linear, None):
-            act_cls_name = activation.__class__.__name__
+            assert isinstance(activation, FunctionType), f"Activation function for layer {layer.name} is not a function"
+            intermediate_tensor_name = f'{output_names[0]}_activation'
+            ret[0]['output_keras_tensor_names'] = (intermediate_tensor_name,)
+            act_cls_name = activation.__name__
             act_config = {
                 'class_name': 'Activation',
                 'activation': act_cls_name,
                 'name': f'{layer.name}_{act_cls_name}',
+                'input_keras_tensor_names': (intermediate_tensor_name,),
+                'output_keras_tensor_names': output_names,
             }
             ret = *ret, act_config
         return ret
@@ -212,19 +222,26 @@ def parse_keras_v3_model(model: 'keras.Model'):
         If a circular dependency is detected.
     """
 
+    assert model.built, "Model must be built before parsing"
+
+    import keras
+
+    if isinstance(model, keras.Sequential):
+        model = model._functional  # everything is functional under the hood lol
+
     from .keras_to_hls import layer_handlers as v2_layer_handlers  # Delayed import to avoid circular import
 
     keras_v3_dispatcher = KerasV3HandlerDispatcher(v3_layer_handlers, v2_layer_handlers)
 
     model_inputs, model_outputs, dependency, tensors = resolve_dependency_relation(model)
 
     satisfied = set()
-    total = len(tensors)
 
     unique_name = UniqueName()
 
     layer_list: list[dict[str, Any]] = []
-    while len(satisfied) < total:
+
+    while any(t not in satisfied for t in model_outputs):
         # Until all tensors in the model are satisfied
         for i, (layer_name, in_tensor_names, out_tensor_names) in enumerate(dependency):
             if not all(t in satisfied for t in in_tensor_names):
@@ -237,13 +254,10 @@ def parse_keras_v3_model(model: 'keras.Model'):
             out_tensors = [tensors[t] for t in out_tensor_names]
 
             _configs = keras_v3_dispatcher(layer, inp_tensors, out_tensors)
-            # Dispatch to v3 handler if available, else fallback to v2
-            # handler
+            # Dispatch to v3 handler if available, else fallback to v2 handler
 
-            # Prevent name conflicts. If a layer is used multiple times,
-            # add a suffix to the name At this stage, connections
-            # between modules are recorded by i/o keras tensor names
-            # (guaranteed unique), thus we can safely rename the layers
+            # Prevent name conflicts. If a layer is used multiple times, add a suffix to the name.
+            # At this stage connections between modules are recorded by i/o keras tensor names
             for _conf in _configs:
                 _conf['name'] = unique_name(_conf['name'])
 

hls4ml/utils/config.py

@@ -1,6 +1,7 @@
 import json
 
 import hls4ml
+import hls4ml.converters.keras_v3_to_hls
 
 
 def create_config(output_dir='my-hls-test', project_name='myproject', backend='Vivado', version='1.0.0', **kwargs):
@@ -157,12 +158,17 @@ def config_from_keras_model(
 
     if isinstance(model, dict):
         model_arch = model
+        reader = hls4ml.converters.KerasModelReader(model)
+        layer_list, _, _, _ = hls4ml.converters.parse_keras_model(model_arch, reader)
     else:
-        model_arch = json.loads(model.to_json())
+        import keras
 
-    reader = hls4ml.converters.KerasModelReader(model)
-
-    layer_list, _, _, _ = hls4ml.converters.parse_keras_model(model_arch, reader)
+        if keras.__version__ > '3.0':
+            layer_list, *_ = hls4ml.converters.parse_keras_v3_model(model)
+        else:
+            model_arch = json.loads(model.to_json())
+            reader = hls4ml.converters.KerasModelReader(model)
+            layer_list, _, _, _ = hls4ml.converters.parse_keras_model(model_arch, reader)
 
     def make_layer_config(layer):
         cls_name = layer['class_name']