Add tests for Mean Dim operator

Author: roman-janik-nxp

backends/nxp/tests/executors.py

@@ -159,15 +159,31 @@ def preprocess(self, data: np.ndarray | dict[int, numpy.ndarray]):
 
 
 class ToChannelFirstPreprocess(TFLiteIOPreprocess):
+    def __init__(self, dim_0_reduced: bool | dict[int, bool] = False):
+        self.dim_0_reduced = dim_0_reduced
+
     def preprocess(self, data: np.ndarray | dict[int, np.ndarray]):
-        def get_channel_first_permutation(tensor):
-            return create_channels_last_to_channels_first_permutation(len(tensor.shape))
+        def get_channel_first_permutation(tensor, dim_0_reduced):
+            tensor_rank = len(tensor.shape)
+            perm = create_channels_last_to_channels_first_permutation(tensor_rank)
+            if dim_0_reduced and tensor_rank > 1:
+                perm[0], perm[1] = perm[1], perm[0]
+            return perm
+
+        transpose_fn = lambda x, rank: np.transpose(x, get_channel_first_permutation(x, rank))
+        if isinstance(data, np.ndarray) and isinstance(self.dim_0_reduced, bool):
+            preprocessed_data = transpose_fn(data, self.dim_0_reduced)
+
+        elif isinstance(data, dict) and isinstance(self.dim_0_reduced, bool):
+            preprocessed_data = {k: transpose_fn(v, self.dim_0_reduced) for k, v in data.items()}
+
+        elif isinstance(data, dict) and isinstance(self.dim_0_reduced, dict):
+            preprocessed_data = {k: transpose_fn(v, self.dim_0_reduced[k]) for k, v in data.items()}
 
-        transpose_fn = lambda x: np.transpose(x, get_channel_first_permutation(x))
-        if isinstance(data, np.ndarray):
-            preprocessed_data = transpose_fn(data)
         else:
-            preprocessed_data = {k: transpose_fn(v) for k, v in data.items()}
+            raise ValueError("Invalid combination of inputs. Data can be either np.ndarray or dict. If original number "
+                             "of dimension is used, it can be only int for np.ndarray data or dict of ints for dict "
+                             "data with same keys.")
         return preprocessed_data
 
 

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

@@ -0,0 +1,107 @@
+import numpy as np
+import pytest
+import torch
+from torch.export import ExportedProgram
+
+from executorch.backends.nxp.backend.edge_program_converter import EdgeProgramToIRConverter
+from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program
+from executorch.backends.nxp.tests.executors import convert_run_compare, ToChannelFirstPreprocess, \
+    ToChannelLastPreprocess
+from executorch.backends.nxp.tests.models import MeanDimConvModule, MeanDimLinearModule
+
+
+@pytest.fixture(autouse=True)
+def reseed_model_per_test_run():
+    torch.manual_seed(23)
+    np.random.seed(23)
+
+
+@pytest.mark.parametrize("input_shape, dim", [
+    pytest.param((1, 4, 8, 8), (-1, -2), id="Dim -1, -2."),
+])
+def test_mean_dim_conv_quant_conversion(mocker, input_shape, dim, keeepdim=True):
+    model = MeanDimConvModule(dim, keeepdim)
+
+    converter_spy = mocker.spy(EdgeProgramToIRConverter, "convert_program")
+
+    # Run conversion
+    _ = to_quantized_edge_program(model, input_shape)
+
+    # Capture generated model
+    tflite_flatbuffers_model, io_formats = converter_spy.spy_return
+
+    # 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)
+
+    convert_run_compare(exported_program, tflite_input_preprocess=ToChannelLastPreprocess(), input_data=input_data,
+                        tflite_output_preprocess=ToChannelFirstPreprocess(), tfl_model=tflite_flatbuffers_model)
+
+
+@pytest.mark.parametrize("input_shape, dim", [
+    pytest.param((1, 32), 0, id="Dim 0."),
+    pytest.param((1, 32), 1, id="Dim 1."),
+])
+@pytest.mark.parametrize("keeepdim", [
+    pytest.param(False, id="Don't keep dim."),
+    pytest.param(True, id="Keep dim."),
+])
+def test_mean_dim_linear_unsupported_quant_conversion(mocker, input_shape, dim, keeepdim):
+    model = MeanDimLinearModule(dim, keeepdim)
+
+    converter_spy = mocker.spy(EdgeProgramToIRConverter, "convert_program")
+
+    # Run conversion
+    edge_program = to_quantized_edge_program(model, input_shape).exported_program()
+    nodes = list(edge_program.graph.nodes)
+
+    # Last 2 dimensions are not used or keepdim is False, cannot be converted to MeanDim, node is not delegated
+    assert nodes[6].target.__name__ == 'aten.mean.dim'
+
+    # Capture generated model
+    tflite_flatbuffers_model, io_formats = converter_spy.spy_return
+
+    # 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)
+
+    convert_run_compare(exported_program, tfl_model=tflite_flatbuffers_model, input_data=input_data)
+
+
+@pytest.mark.parametrize("input_shape, dim", [
+    pytest.param((1, 4, 8, 8), 0, id="Dim 0."),
+    pytest.param((1, 4, 8, 8), 2, id="Dim 2."),
+    pytest.param((1, 4, 8, 8), -1, id="Dim -1."),
+    pytest.param((1, 4, 8, 8), -2, id="Dim -2."),
+    pytest.param((1, 4, 8, 8), (0, 1), id="Dim 0, 1."),
+    pytest.param((1, 4, 8, 8), (1, 3), id="Dim 1, 3."),
+    pytest.param((1, 4, 8, 8), (-1, -3), id="Dim -1, -3."),
+])
+@pytest.mark.parametrize("keeepdim", [
+    pytest.param(False, id="Don't keep dim."),
+    pytest.param(True, id="Keep dim."),
+])
+def test_mean_dim_conv_unsupported_quant_conversion(mocker, input_shape, dim, keeepdim):
+    model = MeanDimConvModule(dim, keeepdim)
+
+    converter_spy = mocker.spy(EdgeProgramToIRConverter, "convert_program")
+
+    # Run conversion
+    edge_program = to_quantized_edge_program(model, input_shape).exported_program()
+    nodes = list(edge_program.graph.nodes)
+
+    # Last 2 dimensions are not used or keepdim is False, cannot be converted to MeanDim, node is not delegated
+    assert nodes[6].target.__name__ == 'aten.mean.dim'
+
+    # Capture generated model
+    tflite_flatbuffers_model, io_formats = converter_spy.spy_return
+
+    # 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)
+
+    convert_run_compare(exported_program, tflite_input_preprocess=ToChannelLastPreprocess(), input_data=input_data,
+                        tflite_output_preprocess=ToChannelFirstPreprocess(), tfl_model=tflite_flatbuffers_model)

backends/nxp/tests/models.py

@@ -249,3 +249,27 @@ def __init__(self):
     @staticmethod
     def forward(x):
         return x + x
+
+
+class MeanDimLinearModule(torch.nn.Module):
+    def __init__(self, dim, keepdim):
+        super().__init__()
+        self.dim = dim
+        self.keepdim = keepdim
+        self.linear = torch.nn.Linear(32, 16)
+
+    def forward(self, x):
+        x = self.linear(x)
+        return torch.mean(x, dim=self.dim, keepdim=self.keepdim)
+
+
+class MeanDimConvModule(torch.nn.Module):
+    def __init__(self, dim, keepdim):
+        super().__init__()
+        self.conv = Conv2dModule(stride=1, padding=1)
+        self.dim = dim
+        self.keepdim = keepdim
+
+    def forward(self, x):
+        x = self.conv(x)
+        return torch.mean(x, dim=self.dim, keepdim=self.keepdim)