NXP backend: Update broadcast condition for broadcastable ops

Author: roman-janik-nxp

backends/nxp/backend/ir/converter/node_converter.py

@@ -5,6 +5,7 @@
 
 import operator
 from abc import ABC, abstractmethod
+from math import prod
 from typing import Callable
 
 import torch
@@ -411,30 +412,31 @@ def uses_shape_broadcasting(node: Node) -> bool:
         )
 
     @staticmethod
-    def at_least_one_input_shape_matches_the_output_shape(node: Node) -> bool:
-        """Determine if given PyTorch fx Node uses at least one input shape broadcasting for it's input nodes or not.
+    def inputs_satisfy_broadcast_condition(node: Node) -> bool:
+        """Determine if given PyTorch fx Node has inputs that satisfy broadcasting conditions for Neutron or not.
 
         :param node: PyTorch fx Node with 'all_input_nodes' initialized.
-        :return: True, if at least one input has the same shape as the output node.
+        :return: True, if at least one input has the same number of elements as the output node.
                  False otherwise.
         """
 
         if node.all_input_nodes is None:
             logger.e(
                 logger.Code.INTERNAL_ERROR,
-                "node_converter.at_least_one_input_shape_matches_the_output_shape(): 'all_input_nodes' are None!",
+                "node_converter.inputs_satisfy_broadcast_condition(): 'all_input_nodes' are None!",
             )
 
         if len(node.all_input_nodes) == 0:
             logger.e(
                 logger.Code.INTERNAL_ERROR,
-                "node_converter.at_least_one_input_shape_matches_the_output_shape(): Operator has no inputs!",
+                "node_converter.inputs_satisfy_broadcast_condition(): Operator has no inputs!",
             )
 
         output_shape = node.meta["val"].shape
+        num_elements = prod(output_shape)
 
         return any(
-            input_tensor.meta["val"].shape == output_shape
+            prod(input_tensor.meta["val"].shape) == num_elements
             for input_tensor in node.all_input_nodes
         )
 

backends/nxp/backend/ir/converter/node_converters/ops_converters/add_tensor_converter.py

@@ -26,7 +26,7 @@ def _is_supported_on_target(
         parameters_mapping: dict[str, Parameter],
         custom_delegation_options: CustomDelegationOptions,
     ) -> bool:
-        if not NodeConverter.at_least_one_input_shape_matches_the_output_shape(node):
+        if not NodeConverter.inputs_satisfy_broadcast_condition(node):
             return False
 
         # If one input is in channel first and ranks of input tensors are not equal, we need to add Transposes

backends/nxp/backend/ir/converter/node_converters/ops_converters/mul_tensor_converter.py

@@ -25,7 +25,7 @@ def _is_supported_on_target(
         parameters_mapping: dict[str, Parameter],
         custom_delegation_options: CustomDelegationOptions,
     ) -> bool:
-        if not NodeConverter.at_least_one_input_shape_matches_the_output_shape(node):
+        if not NodeConverter.inputs_satisfy_broadcast_condition(node):
             return False
 
         # If one input is in channel first and ranks of input tensors are not equal, we need to add Transposes

backends/nxp/backend/ir/converter/node_converters/ops_converters/sub_tensor_converter.py

@@ -26,7 +26,7 @@ def _is_supported_on_target(
         parameters_mapping: dict[str, Parameter],
         custom_delegation_options: CustomDelegationOptions,
     ) -> bool:
-        if not NodeConverter.at_least_one_input_shape_matches_the_output_shape(node):
+        if not NodeConverter.inputs_satisfy_broadcast_condition(node):
             return False
 
         # If one input is in channel first and ranks of input tensors are not equal, we need to add Transposes

backends/nxp/tests/ir/converter/node_converter/test_add_tensor_converter.py

@@ -104,6 +104,10 @@ def test__basic_nsys_inference_qat(self, mocker, request):
             pytest.param(
                 [ModelInputSpec((4,)), ModelInputSpec((4, 4))], id="2 inputs 1D + 2D."
             ),
+            pytest.param(
+                [ModelInputSpec((10,)), ModelInputSpec((1, 1))],
+                id="2 inputs 2D, num_elems of input == num_elems of output",
+            ),
         ],
     )
     def test__broadcast(self, mocker, request, input_spec):

backends/nxp/tests/ir/converter/node_converter/test_mul_tensor_converter.py

@@ -90,6 +90,10 @@ def test__basic_nsys_inference_qat(self, mocker, request, x_input_shape):
             pytest.param(
                 [ModelInputSpec((4,)), ModelInputSpec((4, 4))], id="2 inputs 1D+2D."
             ),
+            pytest.param(
+                [ModelInputSpec((10,)), ModelInputSpec((1, 1))],
+                id="2 inputs 2D, num_elems of input == num_elems of output",
+            ),
         ],
     )
     def test__correct_broadcast(self, input_spec, mocker, request):

backends/nxp/tests/ir/converter/node_converter/test_sub_tensor_converter.py

@@ -104,6 +104,10 @@ def test__basic_nsys_inference_qat(self, mocker, request):
                 [ModelInputSpec((5, 3, 4)), ModelInputSpec((1, 3, 1))],
                 id="2 inputs 3D.",
             ),
+            pytest.param(
+                [ModelInputSpec((10,)), ModelInputSpec((1, 1))],
+                id="2 inputs 2D, num_elems of input == num_elems of output",
+            ),
         ],
     )
     def test__broadcast(self, mocker, request, input_spec):