allow io_stream if used as model output

Author: calad0i

hls4ml/backends/fpga/passes/inplace_stream_flatten.py

@@ -11,11 +11,18 @@ class InplaceStreamFlatten(OptimizerPass):
     """
 
     def match(self, node):
-        # Reshape acts as a Flatten layer when the result has 1 dimension
+        # Layers require flatten data can gather it from the stream, no need for repacking.
+        # Reshape acts as a Flatten layer when the result has 1 dimension. Make it a inplace tensor if it happens.
+
+        if node.model.config.get_config_value('IOType') != 'io_stream':
+            return False
         if not (isinstance(node, Reshape) and len(node.get_output_variable().shape) == 1):
-            # Reshape with multiple outputs will be kept as is, or repack cannot handle different shapes
+            # If is not flatten
+            return False
+        if node.name in node.model.outputs:
+            # If used as model output. Output shape shall be preserved in this case.
             return False
-        return node.model.config.get_config_value('IOType') == 'io_stream'
+        return True
 
     def transform(self, model, node):
         outvar = node.get_output_variable()

hls4ml/backends/fpga/passes/repack_stream.py

@@ -49,7 +49,9 @@ class ReshapeStream(OptimizerPass):
 
     def match(self, node):
         # do not run optimizer pass for a flatten layer (1 output dimension)
-        return isinstance(node, Reshape) and len(node.get_output_variable().shape) > 1
+        if not isinstance(node, Reshape):
+            return False
+        return len(node.get_output_variable().shape) > 1 or node.name in node.model.outputs
 
     def transform(self, model, node):
         if model.config.get_config_value('IOType') != 'io_stream':

hls4ml/model/graph.py

@@ -545,6 +545,8 @@ def remove_node(self, node, rewire=True):
         if len(inputs) == 1:
             # Connect inputs -> $outputs
             if node.name in self.outputs:
+                msg = f'Remove leaf node {node.name} will connect its input node {inputs[0]} to output, but it already is.'
+                assert inputs[0] not in self.outputs, msg
                 self.outputs = [inputs[0] if name == node.name else name for name in self.outputs]
 
         if len(outputs) == 1 and len(inputs) == 1:

test/pytest/test_multiout_network.py

@@ -28,7 +28,9 @@ def model_corner_cases():
     out2 = keras.layers.Dense(16, activation='relu')(out1)
     out2 = keras.layers.Add()([out2, in2])
     out3 = keras.layers.Dense(2)(out1)
-    model = keras.models.Model(inputs=[in1, in2], outputs=[out1, out2, out3])
+    out4 = keras.layers.Dense(2)(out2)
+    out4 = keras.layers.Flatten()(out4)
+    model = keras.models.Model(inputs=[in1, in2], outputs=[out1, out2, out3, out4])
     return model
 
 
@@ -76,7 +78,8 @@ def test_multi_output_nn_corner_cases(model_corner_cases, data_corner_cases, bac
        - when an node removal/insertion is triggered internally
     - a reshape in io_parallel, or flatten in io_stream layer's output is used multiple times
        - and as layer output
-       - and  by layer taking multiple inputs
+       - and by layer taking multiple inputs
+    - a Flatten layer outputs to the model output in io_stream
     """
     output_dir = str(test_root_path / f'hls4mlprj_multiout_network_2_{backend}_{io_type}_{strategy}')
     hls_config = {'Model': {'Precision': 'fixed<32,5>', 'ReuseFactor': 1}, 'Strategy': strategy}
@@ -92,3 +95,4 @@ def test_multi_output_nn_corner_cases(model_corner_cases, data_corner_cases, bac
     assert np.allclose(r_hls[0], r_keras[0], atol=1e-5, rtol=0)
     assert np.allclose(r_hls[1], r_keras[1], atol=1e-5, rtol=0)
     assert np.allclose(r_hls[2], r_keras[2], atol=1e-5, rtol=0)
+    assert np.allclose(r_hls[3], r_keras[3], atol=1e-5, rtol=0)