fix: applied feedback from PR

Author: novak-vaclav

backends/nxp/backend/edge_program_converter.py

@@ -30,7 +30,6 @@
     exir_ops.edge.aten._adaptive_avg_pool2d.default: AdaptiveAvgPool2dConverter,  # noqa F405
     exir_ops.edge.aten.addmm.default: AddMMConverter,  # noqa F405
     exir_ops.edge.aten.add.Tensor: AddTensorConverter,  # noqa F405
-    exir_ops.edge.aten.sub.Tensor: SubTensorConverter,  # noqa F405
     exir_ops.edge.aten.avg_pool2d.default: AvgPool2dConverter,  # noqa F405
     exir_ops.edge.aten.cat.default: CatConverter,  # noqa F405
     exir_ops.edge.aten.clone.default: CloneConverter,  # noqa F405
@@ -43,6 +42,7 @@
     exir_ops.edge.aten.permute_copy.default: PermuteCopyConverter,  # noqa F405
     exir_ops.edge.aten.relu.default: ReLUConverter,  # noqa F405
     exir_ops.edge.aten._softmax.default: SoftmaxConverter,  # noqa F405
+    exir_ops.edge.aten.sub.Tensor: SubTensorConverter,  # noqa F405
     exir_ops.edge.aten.tanh.default: TanhConverter,  # noqa F405
     exir_ops.edge.aten.view_copy.default: ViewCopyConverter,  # noqa F405
     exir_ops.edge.aten.sigmoid.default: SigmoidConverter,  # noqa F405

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

@@ -46,14 +46,16 @@ def _is_supported_in_IR(
         if len(node.args) != 2:
             return False
 
+        # The `alpha` attribute can be represented by adding an extra `Mul` operator.
+        #  However, this is not implemented as `alpha` is rarely used.
         if hasattr(node.kwargs, "alpha"):
             return False
 
         return True
 
     # sub.Tensor Node format: (Tensor self, Tensor other, *, Scalar alpha=1)
     def convert(self, node: Node):
-        """Convert 'sub_tensor' operator to TFLite 'sub'."""
+        """Convert 'sub_tensor' operator to NeutronIR 'Sub'."""
         self.assert_convertible(node)
 
         t_op = self._create_tflite_op_with_io_tensors(node)

backends/nxp/neutron_partitioner.py

@@ -194,7 +194,6 @@ def tag_qdq_clusters(self, nodes: List[torch.fx.Node]):
     exir_ops.edge.aten._adaptive_avg_pool2d.default: AdaptiveAvgPool2dConverter,  # noqa F405
     exir_ops.edge.aten.addmm.default: AddMMConverter,  # noqa F405
     exir_ops.edge.aten.add.Tensor: AddTensorConverter,  # noqa F405
-    exir_ops.edge.aten.sub.Tensor: SubTensorConverter,  # noqa F405
     exir_ops.edge.aten.avg_pool2d.default: AvgPool2dConverter,  # noqa F405
     exir_ops.edge.aten.cat.default: CatConverter,  # noqa F405
     exir_ops.edge.aten.clone.default: CloneConverter,  # noqa F405
@@ -207,6 +206,7 @@ def tag_qdq_clusters(self, nodes: List[torch.fx.Node]):
     exir_ops.edge.aten.mm.default: MMConverter,  # noqa F405
     exir_ops.edge.aten.relu.default: ReLUConverter,  # noqa F405
     exir_ops.edge.aten._softmax.default: SoftmaxConverter,  # noqa F405
+    exir_ops.edge.aten.sub.Tensor: SubTensorConverter,  # noqa F405
     exir_ops.edge.aten.tanh.default: TanhConverter,  # noqa F405
     exir_ops.edge.aten.view_copy.default: ViewCopyConverter,  # noqa F405
     exir_ops.edge.aten.sigmoid.default: SigmoidConverter,  # noqa F405

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

@@ -1,3 +1,7 @@
+# Copyright 2025 NXP
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
 import numpy as np
 import pytest
 import torch

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

@@ -16,6 +16,7 @@
     SubTensorModule,
     SubTensorOneInputModule,
 )
+from executorch.exir.dialects._ops import ops as exir_ops
 from torch.export import ExportedProgram
 
 
@@ -48,8 +49,16 @@ def test_sub_tensor_quant_conversion(mocker, input_shape):
     # Capture converted program
     exported_program: ExportedProgram = converter_spy.call_args.args[1]
 
-    input_data = (np.random.random(input_shape).astype(np.float32) * 50).astype(np.int8)
-    input_data = {0: input_data, 1: input_data}
+    input_data_1 = (np.random.random(input_shape).astype(np.float32) * 50).astype(
+        np.int8
+    )
+    input_data_2 = (np.random.random(input_shape).astype(np.float32) * 50).astype(
+        np.int8
+    )
+    input_data = {0: input_data_1, 1: input_data_2}
+
+    nodes = list(exported_program.graph.nodes)
+    assert nodes[4].name == "aten_sub_tensor"
 
     convert_run_compare(
         exported_program, tfl_model=tflite_flatbuffers_model, input_data=input_data
@@ -81,6 +90,9 @@ def test_sub_tensor_one_input_quant_conversion(mocker, input_shape):
 
     input_data = (np.random.random(input_shape).astype(np.float32) * 50).astype(np.int8)
 
+    nodes = list(exported_program.graph.nodes)
+    assert nodes[2].name == "aten_sub_tensor"
+
     convert_run_compare(
         exported_program, tfl_model=tflite_flatbuffers_model, input_data=input_data
     )
@@ -109,6 +121,9 @@ def test_sub_tensor_w_conv_quant_conversion(mocker, input_shape):
 
     input_data = (np.random.random(input_shape).astype(np.float32) * 50).astype(np.int8)
 
+    nodes = list(exported_program.graph.nodes)
+    assert nodes[9].name == "aten_sub_tensor"
+
     convert_run_compare(
         exported_program,
         input_data,
@@ -142,4 +157,6 @@ def test_sub_tensor_broadcasting_unsupported_quant_conversion(
     nodes = list(edge_program.graph.nodes)
 
     # Broadcast is not supported, node is not converted
-    assert nodes[6].target.__name__ == "aten.sub.Tensor"  # Sub Tensor is not delegated.
+    assert (
+        nodes[6].target == exir_ops.edge.aten.sub.Tensor
+    )  # Sub Tensor is not delegated.