Broadcast refactor

Author: lukamac

Deeploy/DeeployTypes.py

@@ -2046,44 +2046,38 @@ def broadcast(self, ctxt: NetworkContext, default_channels_first: bool = True) -
         inputShapes = [ctxt.lookup(node.name).shape for node in self.node.inputs]
         outputShapes = [ctxt.lookup(node.name).shape for node in self.node.outputs]
 
-        if not "channels_first" in self.mapper.parser.operatorRepresentation:
-            channels_first = default_channels_first
-        else:
-            channels_first = self.mapper.parser.operatorRepresentation['channels_first']
+        opRepr = self.mapper.parser.operatorRepresentation
+        channels_first = opRepr.get("channels_first", default_channels_first)
+        newInputShapes, newOutputShapes = self.computeShapes(inputShapes, outputShapes, opRepr, channels_first)
 
-        newInputShapes, newOutputShapes = self.computeShapes(inputShapes, outputShapes,
-                                                             self.mapper.parser.operatorRepresentation, channels_first)
+        for tensor, shape in zip(self.node.inputs + self.node.outputs, newInputShapes + newOutputShapes):
+            buffer = ctxt.lookup(tensor.name)
+            assert isinstance(buffer, VariableBuffer)
 
-        for node, newShape in zip(self.node.inputs + self.node.outputs, newInputShapes + newOutputShapes):
-            if ctxt.is_local(node.name):
-                ctxt.localObjects[node.name].shape = newShape
+            if ctxt.is_local(tensor.name):
+                buffer.shape = shape
                 # Update shape of tensors in onnx graph
-                node.shape = newShape
+                tensor.shape = shape
 
                 # WIESEP: It is possible that the type was not yet set, so we assume some default type
                 # At this state, we assume that all local buffers are float32 type inference is not yet done.
-                if node.dtype is None:
-                    node.dtype = np.float32
+                if tensor.dtype is None:
+                    tensor.dtype = np.float32
 
-            elif ctxt.is_global(node.name):
-                ctxt.globalObjects[node.name].shape = newShape
-                if isinstance(ctxt.globalObjects[node.name], ConstantBuffer):
+            elif ctxt.is_global(tensor.name):
+                buffer.shape = shape
+                if isinstance(buffer, ConstantBuffer):
 
                     # If the number of elements is equal, reshape
-                    if np.prod(ctxt.globalObjects[node.name].values.shape) == np.prod(newShape):
-                        ctxt.globalObjects[node.name].values.reshape(newShape)
+                    if np.prod(buffer.values.shape) == np.prod(shape):
+                        buffer.values.reshape(shape)
                     # The number of elements SHOULD be lower, and we broadcast
                     else:
                         try:
-                            ctxt.globalObjects[node.name].values = np.broadcast_to(ctxt.globalObjects[node.name].values,
-                                                                                   newShape)
-                        except:
-                            raise RuntimeError(
-                                f"Could not broadcast {node.name} from {ctxt.globalObjects[node.name].values.shape} to {newShape}!"
-                            )
-
-            else:
-                raise KeyError(f'Expected node {node.name} to be in context!')
+                            buffer.values = np.broadcast_to(buffer.values, shape)
+                        except ValueError as e:
+                            raise ValueError(
+                                f"Could not broadcast tensor {tensor.name} of node {self.node.name}.") from e
 
         return ctxt