Fix KLLoss layer example

Author: vloncar

contrib/kl_layer/README.md

@@ -10,7 +10,7 @@ It was developed specifically for [AD@L1 CMS paper](https://www.nature.com/artic
 
 # Usage
 
-`test_extensions` function in `kl_layer.py` contains the example of how to use the KL layer.
+`kl_layer.py` contains the example of how to use the KL layer.
 To run do
 
 ```

contrib/kl_layer/kl_layer.h

@@ -20,7 +20,7 @@ struct distance_config {
 
     // Internal info
     static const unsigned table_size = 1024;
-    static constexpr float exp_range = 1024;
+    static constexpr unsigned exp_range = 8;
 };
 
 template <typename CONFIG_T, int N_TABLE> void init_klloss_exp_table(typename CONFIG_T::exp_table_t table_out[N_TABLE]) {

contrib/kl_layer/kl_layer.py

@@ -21,7 +21,8 @@
 
 import hls4ml
 from hls4ml.converters.keras_to_hls import parse_default_keras_layer
-from hls4ml.model.types import FixedPrecisionType, NamedType
+from hls4ml.model.attributes import ConfigurableAttribute, TypeAttribute
+from hls4ml.model.types import FixedPrecisionType, RoundingMode, SaturationMode
 
 
 # Keras implementation of a KL layer
@@ -47,21 +48,23 @@ def _merge_function(self, inputs):
 class HKLLoss(hls4ml.model.layers.Layer):
     '''hls4ml implementation of a KL loss custom layer'''
 
+    _expected_attributes = [
+        ConfigurableAttribute('table_size', default=1024),
+        ConfigurableAttribute('exp_range', default=8),
+        TypeAttribute('accum'),
+        TypeAttribute(
+            'sum',
+            default=FixedPrecisionType(18, 8, rounding_mode=RoundingMode.RND, saturation_mode=SaturationMode.SAT),
+        ),
+        TypeAttribute(
+            'exp_table',
+            default=FixedPrecisionType(18, 8, rounding_mode=RoundingMode.RND, saturation_mode=SaturationMode.SAT),
+        ),
+    ]
+
     def initialize(self):
         self.add_output_variable(shape=[1], dim_names=[f'KL_LOSS_{self.index}'])
 
-        print(self.attributes)
-        if 'sum_t' not in self.attributes:
-            self.set_attr('sum_t', self.get_attr('accum_t'))
-        if 'exp_table_t' not in self.attributes:
-            self.set_attr(
-                'exp_table_t', NamedType(name=self.name + '_exp_table_t', precision=FixedPrecisionType(width=18, integer=8))
-            )
-        if 'table_size' not in self.attributes:
-            self.set_attr('table_size', 1024)
-        if 'exp_range' not in self.attributes:
-            self.set_attr('exp_range', 8)
-
 
 # Templates
 distance_config_template = """struct config{index} : nnet::distance_config {{
@@ -73,8 +76,8 @@ def initialize(self):
     static const unsigned table_size = {table_size};
     static constexpr float exp_range = {exp_range};
 }};\n"""
-distance_function_template = 'nnet::{distance}<{input1_t}, {input2_t}, {output_t}, {config}>({input1}, {input2}, {output});'
-distance_include_list = ['../../../contrib/kl_layer/kl_layer.h']
+distance_function_template = 'nnet::klloss<{input1_t}, {input2_t}, {output_t}, {config}>({input1}, {input2}, {output});'
+distance_include_list = ['nnet_utils/kl_layer.h']
 
 
 class HKLLossConfigTemplate(hls4ml.backends.template.LayerConfigTemplate):
@@ -96,7 +99,6 @@ def __init__(self):
 
     def format(self, node):
         params = {}
-        params['distance'] = 'klloss'
         params['config'] = f'config{node.index}'
         params['input1_t'] = node.get_input_variable(node.inputs[0]).type.name
         params['input2_t'] = node.get_input_variable(node.inputs[1]).type.name
@@ -134,7 +136,7 @@ def main():
     backend.register_template(HKLLossFunctionTemplate)
 
     # Register HLS implementation
-    p = Path('kl_layer.h')
+    p = Path(__file__).parent / 'kl_layer.h'
     backend.register_source(p)
 
     # Test if it works