From 76d19fbe7a1af76f172e26aa810b66429e3d4856 Mon Sep 17 00:00:00 2001 From: Martin Pavella Date: Thu, 11 Jun 2026 11:08:51 +0200 Subject: [PATCH] NXP backend: Resolve `mean.dim` format related issue. --- .../ops_converters/mean_dim_converter.py | 125 ++++++- backends/nxp/backend/node_format_inference.py | 73 ++-- .../node_converter/test_mean_dim_converter.py | 318 +++++++++++++----- backends/nxp/tests/ops_aliases.py | 1 + 4 files changed, 417 insertions(+), 100 deletions(-) diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/mean_dim_converter.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/mean_dim_converter.py index a76abfbef91..8674bf697c7 100644 --- a/backends/nxp/backend/ir/converter/node_converters/ops_converters/mean_dim_converter.py +++ b/backends/nxp/backend/ir/converter/node_converters/ops_converters/mean_dim_converter.py @@ -5,6 +5,9 @@ import torch +from executorch.backends.nxp.backend.data_format import DataFormat +from executorch.backends.nxp.backend.ir.converter.conversion import translator +from executorch.backends.nxp.backend.ir.converter.conversion.common import OpsList from executorch.backends.nxp.backend.ir.converter.conversion.translator import ( create_channels_last_to_channels_first_permutation, ) @@ -89,10 +92,15 @@ def _is_supported_in_IR( def _to_pos_dim(d: int, rank: int): return d + rank if d < 0 else d + @staticmethod + def _normalize_dim(dim: list[int], rank: int) -> list[int]: + # convert negative index to positive + return [MeanDimConverter._to_pos_dim(d, rank) for d in dim] + @staticmethod def _normalize_and_to_channel_last_dim(dim: list[int], rank: int) -> list[int]: # convert negative index to positive - dim = [MeanDimConverter._to_pos_dim(d, rank) for d in dim] + dim = MeanDimConverter._normalize_dim(dim, rank) perm = create_channels_last_to_channels_first_permutation(rank, True) dim = [perm[d] for d in dim] @@ -106,6 +114,114 @@ def _get_attrs(node: Node) -> tuple[list[int], bool]: keepdim = node.args[2] if len(node.args) >= 3 else False return dim, keepdim + def _get_dim_and_handle_io_formats( + self, ops: OpsList, dim: list[int], keep_dim: bool + ): + t_op = ops.middle_op + x = t_op.tmp_inputs[0] + y = t_op.tmp_outputs[0] + + channels_last_input = x.tensor_format.is_channels_last() + channels_last_output = y.tensor_format.is_channels_last() + formatless_input = not channels_last_input + formatless_output = not channels_last_output + + dim = self._normalize_dim(dim, x.rank) + + if keep_dim: + # The rank is preserved and the io formats should always be equal. + assert ( + x.tensor_format == y.tensor_format + ), "NXP backend: There is a bug in `mean.dim` format inference." + + # Just adjust the dim to match the input format. + if channels_last_input: + dim = self._normalize_and_to_channel_last_dim(dim, x.rank) + + else: + # `keep_dim = False`, so the output rank != input rank, and the operator changes the tensor format. + + if channels_last_input and formatless_output: + if 1 in dim: + # If we are reducing over the channels, the channels dimension gets removed and the output ends up + # exactly equal in channels last and channels first, regardless of which other dimensions are + # removed. Therefore, we can just adjust the `dim` and we don't need to insert any `Transpose` ops. + dim = self._normalize_and_to_channel_last_dim(dim, x.rank) + elif all(spatial_dim in dim for spatial_dim in range(2, x.rank)): + # All spatial dims are reduced, leaving only batch and channels (both optionally). So the result is + # equal in channels first and channels last as long as we adjust the `dim` to match a channels last + # input (similarly to the case above). + dim = self._normalize_and_to_channel_last_dim(dim, x.rank) + else: + # If the channels dimension is preserved, we must transpose the input to channels first (to match + # the edge model) and we must keep the `dim` unchanged (referencing channels first dimensions). + # Otherwise, the output would not match the input. + to_channels_first_perm = ( + translator.create_channels_last_to_channels_first_permutation( + x.rank + ) + ) + ops.add_pre( + self.builder.create_transpose_operator_before( + t_op, 0, to_channels_first_perm + ) + ) + t_op.tmp_inputs[0].tensor_format = DataFormat.CHANNELS_FIRST + + elif formatless_input and channels_last_output: + # We need apply the `mean` with the original `dim`, which will produce a channels first output. Then, + # we need to append a `Transpose` operator to make the output channels last. + to_channels_last_perm = ( + translator.create_channels_first_to_channels_last_permutation( + y.rank, True + ) + ) + ops.add_post( + self.builder.create_transpose_operator_after( + t_op, 0, to_channels_last_perm + ) + ) + t_op.tmp_outputs[0].tensor_format = DataFormat.CHANNELS_FIRST + + elif formatless_input and formatless_output: + # No action needed. + pass + + else: # channels_last_input and channels_last_output + # This case cannot currently occur, as it would require the case: + # channels last 4D -> mean -> channels_last 3D + # which cannot currently happen as the 3D conv/pooling/... is supported by adding `view_copy` nodes in + # the edge dialect and converting the node to 4D, and the `view_copy` nodes prevent the propagation of + # the format to the `mean.dim` output. + # Therefore, the implementation cannot be tested. But from experience with other operators, it should + # work correctly. We just need to add 2 `Transpose` ops to make the IO channels first, and keep the + # `dim` unchanged. + to_channels_first_perm = ( + translator.create_channels_last_to_channels_first_permutation( + x.rank + ) + ) + ops.add_pre( + self.builder.create_transpose_operator_before( + t_op, 0, to_channels_first_perm + ) + ) + t_op.tmp_inputs[0].tensor_format = DataFormat.CHANNELS_FIRST + + to_channels_last_perm = ( + translator.create_channels_first_to_channels_last_permutation( + y.rank, True + ) + ) + ops.add_post( + self.builder.create_transpose_operator_after( + t_op, 0, to_channels_last_perm + ) + ) + t_op.tmp_outputs[0].tensor_format = DataFormat.CHANNELS_FIRST + + return dim + def convert(self, node: Node): """Convert the 'mean.dim' operator to NeutronIR 'Mean'. The ExecuTorch schema is: @@ -123,10 +239,9 @@ def convert(self, node: Node): t_op = self._create_tflite_op_with_io_tensors(node) t_op.builtin_options = mean_options.Mean(keepdim) - x = t_op.tmp_inputs[0] - if x.tensor_format.is_channels_last(): - dim = self._normalize_and_to_channel_last_dim(dim, x.rank) + ops = OpsList(middle_op=t_op) + dim = self._get_dim_and_handle_io_formats(ops, dim, keepdim) convert_axes_from_attribute(t_op, self.builder, dim) - self.builder.append_operators([t_op]) + self.builder.append_operators(ops.flatten()) diff --git a/backends/nxp/backend/node_format_inference.py b/backends/nxp/backend/node_format_inference.py index 65e34b7fbde..030873c88ab 100644 --- a/backends/nxp/backend/node_format_inference.py +++ b/backends/nxp/backend/node_format_inference.py @@ -9,10 +9,27 @@ import torch from executorch.backends.nxp.backend.data_format import DataFormat, NXP_NODE_FORMAT - -from executorch.backends.nxp.backend.edge_helper import is_channels_last_dim_order +from executorch.backends.nxp.backend.edge_helper import ( + is_channels_last_dim_order, + try_get_arg, +) from executorch.backends.nxp.backend.edge_program_converter import functions_converters -from executorch.exir.dialects._ops import ops as exir_ops +from executorch.backends.nxp.tests.ops_aliases import ( + AdaptiveAvgPool2D, + AvgPool2D, + Convolution, + DequantizePerChannel, + DequantizePerTensor, + GetItem, + MaxPool2D, + MaxPool2DWithIndices, + MeanDim, + PermuteCopy, + QuantizePerTensor, + UpsampleBilinear2D, + UpsampleNearest2D, + ViewCopy, +) from executorch.exir.dialects.edge._ops import EdgeOpOverload from torch.export import ExportedProgram from torch.fx import Node @@ -25,21 +42,22 @@ class NodeFormatInference: # The op in the dictionary is mapped to a dictionary, which holds indices to input nodes # that are always channels first. ops_with_channels_first_nodes = { - exir_ops.edge.aten._adaptive_avg_pool2d.default: {"inputs": [0]}, + AdaptiveAvgPool2D: {"inputs": [0]}, torch.ops.aten.adaptive_avg_pool2d.default: {"inputs": [0]}, - exir_ops.edge.aten.avg_pool2d.default: {"inputs": [0]}, - exir_ops.edge.aten.convolution.default: {"inputs": [0, 1]}, - exir_ops.edge.aten.max_pool2d_with_indices.default: {"inputs": [0]}, - exir_ops.edge.aten.max_pool2d.default: {"inputs": [0]}, - exir_ops.edge.aten.upsample_bilinear2d.vec: {"inputs": [0]}, - exir_ops.edge.aten.upsample_nearest2d.vec: {"inputs": [0]}, + AvgPool2D: {"inputs": [0]}, + Convolution: {"inputs": [0, 1]}, + MaxPool2DWithIndices: {"inputs": [0]}, + MaxPool2D: {"inputs": [0]}, + UpsampleBilinear2D: {"inputs": [0]}, + UpsampleNearest2D: {"inputs": [0]}, } # A set of Edge Aten ops, which have the ability to change the format (for example - input nodes # are channels first but output is formatless). ops_that_can_change_tensor_format = { - exir_ops.edge.aten.view_copy.default, - exir_ops.edge.aten.permute_copy.default, + ViewCopy, + PermuteCopy, + MeanDim, } _type_changed_during_last_run: bool @@ -71,10 +89,10 @@ def __init__(self, edge_program: ExportedProgram, only_for_op_support_check=Fals self._type_changed_during_last_run = False self._known_targets = list(functions_converters) + [ - exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, - exir_ops.edge.quantized_decomposed.dequantize_per_channel.default, - exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, - operator.getitem, + DequantizePerTensor, + DequantizePerChannel, + QuantizePerTensor, + GetItem, ] def identify_node_formats(self): @@ -104,10 +122,7 @@ def _infer_format_of_nodes(self, node: Node): self._handle_node_which_uses_channels_first_format(node) elif op_type in self.ops_that_can_change_tensor_format: - if op_type in [ - exir_ops.edge.aten.view_copy.default, - exir_ops.edge.aten.permute_copy.default, - ]: + if op_type in [ViewCopy, PermuteCopy]: # Try to assign the `formatless` format to the input and output. The converter will then handle the # transition. # Note: If the format for the input/output has already been assigned as channels first, it will NOT be @@ -119,10 +134,28 @@ def _infer_format_of_nodes(self, node: Node): self._node_inputs[node][0], DataFormat.FORMATLESS ) + elif op_type == MeanDim: + # The operator schema is: + # mean.dim(Tensor self, int[1]? dim, bool keepdim=False, *, ScalarType? dtype=None) -> Tensor + keep_dim = try_get_arg(node, 2) or False + if keep_dim: + # The operator preserves the rank, so we can handle it as an operator that can use any node format. + self._handle_node_which_can_use_any_node_format(node) + else: + # The operator removes dimensions, so the IO must be marked as `formatless` (unless overridden by + # channels first of course). + self._assign_format_to_node( + self._node_outputs[node][0], DataFormat.FORMATLESS + ) + self._assign_format_to_node( + self._node_inputs[node][0], DataFormat.FORMATLESS + ) + else: logger.error( f"Node format inference for node type: {op_type} not found!" ) + elif node.op != "call_function" or ( hasattr(node, "target") and node.target in self._known_targets ): diff --git a/backends/nxp/tests/ir/converter/node_converter/test_mean_dim_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_mean_dim_converter.py index 8195581c0f6..a003d11ef59 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_mean_dim_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_mean_dim_converter.py @@ -9,6 +9,18 @@ import pytest import torch +from executorch.backends.nxp.backend.ir.converter.builder.model_builder import ( + ModelBuilder, +) +from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options.max_pool_2d_options import ( + MaxPool2D, +) +from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options.mean_options import ( + Mean, +) +from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options.transpose_options import ( + Transpose, +) from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops @@ -50,62 +62,70 @@ def forward(self, x): class MaxPoolMeanDimModule(torch.nn.Module): + @staticmethod + def noop_max_pool_2d(x): + """Call `torch.max_pool2d` that is a NoOp, but it enforces the ChannelsFirst format in the `NodeFormatInference`.""" + return torch.max_pool2d(x, kernel_size=1) + def __init__(self, dim, keepdim): super().__init__() self.dim, self.keepdim = dim, keepdim def forward(self, x): - x = torch.max_pool2d( - x, kernel_size=1 - ) # NoOp, but it enforces the channels first format. - return torch.mean(x, dim=self.dim, keepdim=self.keepdim) + x = self.noop_max_pool_2d(x) + x = torch.mean(x, dim=self.dim, keepdim=self.keepdim) + return x -class TestMeanDim: +class MeanDimMaxPoolModule(MaxPoolMeanDimModule): + def forward(self, x): + x = torch.mean(x, dim=self.dim, keepdim=self.keepdim) + x = self.noop_max_pool_2d(x) + return x + + +def assert_delegated( + model, + input_shape, + mocker, + use_qat=False, + expected_delegated_ops=None, +): + if expected_delegated_ops is None: + expected_delegated_ops = {MeanDim: 1} + + graph_verifier = DetailedGraphVerifier( + mocker, + expected_delegated_ops=expected_delegated_ops, + expected_non_delegated_ops={}, + ) - # noinspection PyMethodMayBeStatic - def assert_delegated( - self, + # Cover also negative values to thoroughly test the operator. + dataset_creator = RandomDatasetCreator(low=-2, high=2) + + remove_quant_io_ops = True # Use quantized dataset. + output_comparator = AllCloseOutputComparator(atol=1) # Allow single bit error. + + lower_run_compare( model, input_shape, - mocker, - use_qat=False, - atol=None, - expected_delegated_ops=None, - ): - if expected_delegated_ops is None: - expected_delegated_ops = {MeanDim: 1} - - graph_verifier = DetailedGraphVerifier( - mocker, - expected_delegated_ops=expected_delegated_ops, - expected_non_delegated_ops={}, - ) + graph_verifier, + dataset_creator, + output_comparator, + use_qat=use_qat, + remove_quant_io_ops=remove_quant_io_ops, + ) - # Cover also negative values to thoroughly test the operator. - dataset_creator = RandomDatasetCreator(low=-2, high=2) - kwargs = {"atol": atol} if atol is not None else {} - output_comparator = AllCloseOutputComparator(**kwargs) +def assert_not_delegated(model, input_shape): + delegated_ep = to_quantized_edge_program(model, input_shape).exported_program() - lower_run_compare( - model, - input_shape, - graph_verifier, - dataset_creator, - output_comparator, - use_qat=use_qat, - ) + # Make sure the `mean` was NOT delegated. + assert not graph_contains_any_of_ops(delegated_ep.graph, [ExecutorchDelegateCall]) + assert graph_contains_any_of_ops(delegated_ep.graph, [MeanDim]) - # noinspection PyMethodMayBeStatic - def assert_not_delegated(self, model, input_shape): - delegated_ep = to_quantized_edge_program(model, input_shape).exported_program() - # Make sure the `mean` was NOT delegated. - assert not graph_contains_any_of_ops( - delegated_ep.graph, [ExecutorchDelegateCall] - ) - assert graph_contains_any_of_ops(delegated_ep.graph, [MeanDim]) +class TestMeanDim: @pytest.fixture(params=[True, False], ids=lambda keep_dim: f"keep_dim = {keep_dim}") def keep_dim(self, request): @@ -114,7 +134,7 @@ def keep_dim(self, request): def test__basic_nsys_inference__qat(self, mocker, use_qat, keep_dim): input_shape = (23,) model = MeanDimModule(0, keep_dim) - self.assert_delegated(model, input_shape, mocker, use_qat=use_qat) + assert_delegated(model, input_shape, mocker, use_qat=use_qat) @pytest.mark.parametrize( "input_shape, dim", @@ -130,10 +150,7 @@ def test__basic_nsys_inference__qat(self, mocker, use_qat, keep_dim): ) def test__single_dims(self, mocker, input_shape, dim, keep_dim): model = MeanDimModule(dim, keep_dim) - # Relatively large error, but it is actually equal to the output scale, so it is a single bit error. - # TODO Replace with quantized dataset testing and `atol = 1`. - atol = 0.014 - self.assert_delegated(model, input_shape, mocker, atol=atol) + assert_delegated(model, input_shape, mocker) @pytest.mark.parametrize( "input_shape, dim", @@ -147,10 +164,7 @@ def test__single_dims(self, mocker, input_shape, dim, keep_dim): ) def test__tuple_dims(self, mocker, input_shape, dim, keep_dim): model = MeanDimModule(dim, keep_dim) - # Relatively large error, but it is actually equal to the output scale, so it is a single bit error. - # TODO Replace with quantized dataset testing and `atol = 1`. - atol = 0.015 - self.assert_delegated(model, input_shape, mocker, atol=atol) + assert_delegated(model, input_shape, mocker) @pytest.mark.parametrize( "input_shape, dim", @@ -162,7 +176,7 @@ def test__tuple_dims(self, mocker, input_shape, dim, keep_dim): def test__noop__only_node__not_delegated(self, input_shape, dim): keep_dim = True # Reduction over a dimension of size `1` with `keep_dim=True` is a no-op. model = MeanDimModule(dim, keep_dim) - self.assert_not_delegated(model, input_shape) + assert_not_delegated(model, input_shape) @pytest.mark.parametrize( "input_shape, dim", @@ -174,7 +188,7 @@ def test__noop__only_node__not_delegated(self, input_shape, dim): def test__noop__not_only_node__delegated(self, mocker, input_shape, dim): keep_dim = True # Reduction over a dimension of size `1` with `keep_dim=True` is a no-op. model = MeanDimAddModule(dim, keep_dim) - self.assert_delegated( + assert_delegated( model, input_shape, mocker, @@ -186,6 +200,7 @@ def test__noop__not_only_node__delegated(self, mocker, input_shape, dim): [ pytest.param((3, 1, 4), 1, id="3D, dim = 1."), pytest.param((3, 1, 4, 1, 5), -2, id="5D, dim = -2."), + pytest.param((1, 7, 3, 3), [0], id="4D, dim = [0]."), ], ) def test__no_reduction__keepdim_false__delegated(self, mocker, input_shape, dim): @@ -194,36 +209,189 @@ def test__no_reduction__keepdim_false__delegated(self, mocker, input_shape, dim) # but with `keep_dim=False` it changes the shape so it's not a noop and is therefore delegated successfully. keep_dim = False model = MeanDimModule(dim, keep_dim) - self.assert_delegated(model, input_shape, mocker) + assert_delegated(model, input_shape, mocker) - @pytest.mark.parametrize( - "input_shape, dim", - [((1, 7, 3, 3), 1)], - ids=lambda val: f"shape={val}" if isinstance(val, tuple) else f"dim={val}", - ) - @pytest.mark.parametrize( - "keep_dim", - [ - pytest.param(True), - pytest.param( - False, - marks=pytest.mark.xfail( - strict=True, reason="Known format inference bug (EIEX-937)." - ), - ), - ], - ids=lambda kd: f"keep_dim={kd}", - ) - def test__channels_first__keep_dim__true(self, mocker, input_shape, dim, keep_dim): + def test__channels_first__keep_dim__true(self, mocker): # Just 1 test case to verify correct handling of the `dim`. # Most cases fall into the single bit error case, and since this test uses 2 operators, the error accumulates # and the final error is larger. We cannot with 100% certainty say that the error is only caused by the single # bit errors and not related to the format. That's why only this 1 case with no errors is used. - - model = MaxPoolMeanDimModule(dim, keep_dim) - self.assert_delegated( + input_shape, dim = (1, 7, 3, 3), 1 + model = MaxPoolMeanDimModule(dim, True) + assert_delegated( model, input_shape, mocker, expected_delegated_ops={MaxPool2DWithIndices: 1, GetItem: 1, MeanDim: 1}, ) + + class TestKeepDimFalseFormatHandling: + """When `keep_dim = False`, the `mean.dim` operator changes the rank, so the format have to be explicitly + handled. The tests in this class focus on the related edge cases. + """ + + def _assert_neutron_ir_model_has_ops( + self, model_builder_finish_spy, expected_ops + ): + assert ( + model_builder_finish_spy.call_count == 1 + ), "Conversion to Neutron IR happened multiple times." + + neutron_ir_ops = model_builder_finish_spy.spy_return.sub_graphs[ + 0 + ].operators.vector + assert len(neutron_ir_ops) == len( + expected_ops + ), "Neutron IR model doesn't have the expected number of ops." + + for op, expected_op in zip(neutron_ir_ops, expected_ops, strict=True): + assert isinstance( + op.builtin_options, expected_op + ), f"Expected {expected_op}, got {op}." + + @pytest.mark.parametrize( + "dim", + [ + 1, + [0, -3], + (-4, 1, 2), + [-3, 3], + [1, 2, 3], + ], + ids=lambda dim: f"dim={dim}", + ) + def test__channels_first_input__reducing_channels(self, mocker, dim): + # If the channels dimension is reduced (removed), the `mean` output will always be equal in channels first + # and channels last, so no `Transpose` ops are added. + input_shape = (1, 7, 3, 3) + model = MaxPoolMeanDimModule(dim, False) + + model_builder_finish_spy = mocker.spy(ModelBuilder, "finish") + assert_delegated( + model, + input_shape, + mocker, + expected_delegated_ops={ + MaxPool2DWithIndices: 1, + GetItem: 1, + MeanDim: 1, + }, + ) + self._assert_neutron_ir_model_has_ops( + model_builder_finish_spy, + expected_ops=[ + Transpose, + MaxPool2D, + Mean, + ], + ) + + @pytest.mark.parametrize( + "dim", + [ + (2, 3), + [1, -2, 3], + [-1, -2, 0], + ], + ids=lambda dim: f"dim={dim}", + ) + def test__channels_first_input__reducing_all_spatial_dims(self, mocker, dim): + # If tall he spatial dimensions are reduced (removed), the `mean` output will always be equal in channels + # first and channels last, so no `Transpose` ops are added. + input_shape = (1, 7, 3, 3) + model = MaxPoolMeanDimModule(dim, False) + + model_builder_finish_spy = mocker.spy(ModelBuilder, "finish") + assert_delegated( + model, + input_shape, + mocker, + expected_delegated_ops={ + MaxPool2DWithIndices: 1, + GetItem: 1, + MeanDim: 1, + }, + ) + self._assert_neutron_ir_model_has_ops( + model_builder_finish_spy, + expected_ops=[ + Transpose, + MaxPool2D, + Mean, + ], + ) + + @pytest.mark.xfail(strict=True, reason="Known Neutron bug (AIR-14726).") + @pytest.mark.parametrize( + "dim", + [ + 0, + (2,), + [-1, 0], + ], + ids=lambda dim: f"dim={dim}", + ) + def test__channels_first_input__not_reducing_channels_or_all_spatial_dims( + self, mocker, dim + ): + # If the channels dimension is not reduced, a `Transpose` operator must be added to make the input channels + # first in Neutron IR. + + input_shape = (1, 7, 3, 3) + model = MaxPoolMeanDimModule(dim, False) + + model_builder_finish_spy = mocker.spy(ModelBuilder, "finish") + assert_delegated( + model, + input_shape, + mocker, + expected_delegated_ops={ + MaxPool2DWithIndices: 1, + GetItem: 1, + MeanDim: 1, + }, + ) + + self._assert_neutron_ir_model_has_ops( + model_builder_finish_spy, + expected_ops=[ + Transpose, + MaxPool2D, + Transpose, # The necessary `Transpose` operator. + Mean, + ], + ) + + @pytest.mark.parametrize( + "input_shape, dim", + [ + pytest.param((2, 3, 4, 5, 6), 0, id="dim=0, 5D->4D"), + pytest.param((2, 3, 4, 5, 6), [-3], id="dim=[-3], 5D->4D"), + pytest.param((1, 2, 3, 4, 5, 6), (1, -1), id="dim=(1, -1), 6D->4D"), + ], + ids=lambda dim: f"dim={dim}", + ) + def test__channels_first_output(self, mocker, input_shape, dim): + model = MeanDimMaxPoolModule(dim, False) + + model_builder_finish_spy = mocker.spy(ModelBuilder, "finish") + assert_delegated( + model, + input_shape, + mocker, + expected_delegated_ops={ + MaxPool2DWithIndices: 1, + GetItem: 1, + MeanDim: 1, + }, + ) + + self._assert_neutron_ir_model_has_ops( + model_builder_finish_spy, + expected_ops=[ + Mean, + Transpose, # The necessary `Transpose` operator. + MaxPool2D, + Transpose, + ], + ) diff --git a/backends/nxp/tests/ops_aliases.py b/backends/nxp/tests/ops_aliases.py index 92f3193b19a..dd39777526f 100644 --- a/backends/nxp/tests/ops_aliases.py +++ b/backends/nxp/tests/ops_aliases.py @@ -29,6 +29,7 @@ HardTanh = exir_ops.edge.aten.hardtanh.default HardTanh_ = exir_ops.edge.aten.hardtanh_.default LeakyRelu = exir_ops.edge.aten.leaky_relu.default +MaxPool2D = exir_ops.edge.aten.max_pool2d.default MaxPool2DWithIndices = exir_ops.edge.aten.max_pool2d_with_indices.default MeanDim = exir_ops.edge.aten.mean.dim MulTensor = exir_ops.edge.aten.mul.Tensor