From 603dfe437c9f40c311281d03982f8746abd775dc Mon Sep 17 00:00:00 2001
From: Martin Pavella <martin.pavella@nxp.com>
Date: Thu, 23 Apr 2026 15:30:19 +0200
Subject: [PATCH 1/4] Unify `to_quantized_edge_program` and
 `to_quantized_executorch_program` duplicates.

---
 backends/nxp/quantizer/utils.py               |  23 +-
 backends/nxp/tests/executorch_pipeline.py     | 157 +++++++++----
 backends/nxp/tests_models/dataset_creator.py  |   2 +-
 backends/nxp/tests_models/executors.py        |  21 +-
 backends/nxp/tests_models/model_input_spec.py |  18 --
 backends/nxp/tests_models/test_cifarnet.py    |   2 +-
 backends/nxp/tests_models/utils.py            | 213 +-----------------
 7 files changed, 151 insertions(+), 285 deletions(-)
 delete mode 100644 backends/nxp/tests_models/model_input_spec.py

diff --git a/backends/nxp/quantizer/utils.py b/backends/nxp/quantizer/utils.py
index da2448fb773..09ee0e5b2f8 100644
--- a/backends/nxp/quantizer/utils.py
+++ b/backends/nxp/quantizer/utils.py
@@ -10,7 +10,7 @@
 import itertools
 from collections import OrderedDict
 from collections.abc import Iterable
-from typing import Any, Dict, List, Tuple, Type
+from typing import Any, Callable, Dict, List, Tuple, Type
 
 import torch
 from executorch.backends.nxp.aten_passes.fuse_batch_norm_with_linear_pass import (
@@ -30,8 +30,10 @@
     check_subgraphs_connected,
     SourcePartition,
 )
+
 from torchao.quantization.pt2e import (
     move_exported_model_to_eval,
+    move_exported_model_to_train,
     ObserverOrFakeQuantize,
 )
 from torchao.quantization.pt2e.quantize_pt2e import (
@@ -176,16 +178,17 @@ def calibrate_and_quantize(
     calibration_inputs: Iterable[tuple[torch.Tensor, ...]],
     quantizer: Quantizer,
     is_qat: bool = False,
+    train_fn: Callable[[torch.fx.GraphModule], None] | None = None,
 ) -> fx.GraphModule:
     """Quantize the provided model.
 
     :param model: Aten model (or it's GraphModule representation) to quantize.
-    :param calibration_inputs: Either a tuple of calibration input tensors where each element corresponds to a model
-                                input. Or an iterator over such tuples.
+    :param calibration_inputs: An iterator over tuples of calibration input tensors where each tensor corresponds to a
+                                model input.
     :param quantizer: Quantizer to use.
     :param is_qat: Whether quantization is done using Quantization Aware Training (QAT) or not.
                     Note: In QAT mode, training is not performed. Only calibration (in eval mode) is done.
-
+    :param train_fn: Optional training function to be called during QAT.
     :return: Quantized GraphModule.
     """
 
@@ -195,12 +198,20 @@ def calibrate_and_quantize(
     if is_qat:
         m = prepare_qat_pt2e(model, quantizer)
         m = AddSimulatedLinearBatchNormFusionQATPass()(m).graph_module
+
+        if train_fn:
+            m = move_exported_model_to_train(m)
+            train_fn(m)
+
         m = move_exported_model_to_eval(m)
+        m = RemoveSimulatedLinearBatchNormFusionQATPass()(m).graph_module
+        m = FuseBatchNormWithLinearPass()(m).graph_module
     else:
         m = prepare_pt2e(model, quantizer)
 
-    for data in calibration_inputs:
-        m(*data)
+    if not is_qat or (is_qat and not train_fn):
+        for data in calibration_inputs:
+            m(*data)
 
     if is_qat:
         m = RemoveSimulatedLinearBatchNormFusionQATPass()(m).graph_module
diff --git a/backends/nxp/tests/executorch_pipeline.py b/backends/nxp/tests/executorch_pipeline.py
index bfe7aca0e27..55c8b1e2b5f 100644
--- a/backends/nxp/tests/executorch_pipeline.py
+++ b/backends/nxp/tests/executorch_pipeline.py
@@ -8,9 +8,10 @@
 import re
 from dataclasses import dataclass
 from functools import partial
-from typing import Callable
+from typing import Callable, Iterable
 
 import eiq_neutron_sdk
+import numpy as np
 import torch
 
 from executorch import exir
@@ -28,7 +29,6 @@
     RemoveIOQuantOpsPass,
 )
 from executorch.backends.nxp.neutron_partitioner import NeutronPartitioner
-
 from executorch.backends.nxp.nxp_backend import (
     core_aten_ops_exception_list,
     generate_neutron_compile_spec,
@@ -42,7 +42,7 @@
     ExecutorchProgramManager,
     to_edge_transform_and_lower,
 )
-from torch import nn
+from torch import memory_format, nn
 from torch.export import export
 from torchao.quantization.pt2e.quantizer import Quantizer
 
@@ -52,7 +52,9 @@
 @dataclass
 class ModelInputSpec:
     shape: tuple[int, ...]
+    type: np.dtype = np.float32
     dtype: torch.dtype = torch.float32
+    dim_order: memory_format = torch.contiguous_format
 
 
 def handle_kernel_selection(model_name: str = ""):
@@ -81,11 +83,11 @@ def handle_kernel_selection(model_name: str = ""):
 
 
 def get_random_calibration_inputs(
-    input_spec: tuple[ModelInputSpec, ...]
+    input_spec: Iterable[ModelInputSpec], num_samples: int = 4
 ) -> list[tuple[torch.Tensor, ...]]:
     return [
         tuple([torch.randn(spec.shape, dtype=spec.dtype) for spec in input_spec])
-        for _ in range(4)
+        for _ in range(num_samples)
     ]
 
 
@@ -94,35 +96,91 @@ def _get_default_quantizer(target_spec: NeutronTargetSpec, use_qat: bool) -> Qua
 
 
 def to_model_input_spec(
-    input_spec: tuple[ModelInputSpec, ...] | tuple[int, ...] | list[tuple[int, ...]]
+    input_spec: Iterable[ModelInputSpec] | tuple[int, ...] | list[tuple[int, ...]]
 ) -> tuple[ModelInputSpec, ...]:
-    if isinstance(input_spec, tuple) and all(
-        isinstance(spec, ModelInputSpec) for spec in input_spec
-    ):
-        return input_spec
-
-    elif isinstance(input_spec, tuple) and all(
-        isinstance(spec, int) for spec in input_spec
-    ):
-        return (ModelInputSpec(input_spec),)
-
-    elif isinstance(input_spec, list) and all(
-        isinstance(input_shape, tuple) for input_shape in input_spec
-    ):
-        return tuple([ModelInputSpec(spec) for spec in input_spec])
-    else:
-        raise TypeError(f"Unsupported type {type(input_spec)}")
+    match input_spec:
+        case _ if isinstance(input_spec, Iterable) and all(
+            isinstance(spec, ModelInputSpec) for spec in input_spec
+        ):
+            return tuple(input_spec)
+        case tuple() if all(isinstance(spec, int) for spec in input_spec):
+            return (ModelInputSpec(input_spec),)
+        case list() if all(
+            isinstance(input_shape, tuple) for input_shape in input_spec
+        ):
+            return tuple([ModelInputSpec(spec) for spec in input_spec])
+        case _:
+            raise TypeError(f"Unsupported type {type(input_spec)}")
+
+
+GetCalibrationInputsFn = Callable[
+    [tuple[ModelInputSpec, ...]], Iterable[tuple[torch.Tensor, ...]]
+]
+
+
+def get_calibration_inputs_fn_from_dataset_dir(dataset_dir) -> GetCalibrationInputsFn:
+    def _nested(
+        input_spec: tuple[ModelInputSpec, ...]
+    ) -> Iterable[tuple[torch.Tensor, ...]]:
+        data = sorted(os.listdir(dataset_dir))
+        inputs_needed = len(input_spec)
+
+        for path in data:
+            path = os.path.join(dataset_dir, path)
+            files = []
+
+            if os.path.isdir(path):
+                files = [os.path.join(path, x) for x in sorted(os.listdir(path))]
+            else:
+                files.append(path)
+
+            input_data = []
+            for idx, file in enumerate(files):
+                if len(input_data) == inputs_needed:
+                    break
+
+                tensor = np.fromfile(file, dtype=input_spec[idx].type).reshape(
+                    input_spec[idx].shape
+                )
+                input_data += (torch.from_numpy(tensor),)
+                continue
+
+            if len(input_data) < inputs_needed:
+                continue
+
+            yield tuple(input_data)
+
+    return _nested
+
+
+def _get_example_input(
+    input_spec: tuple[ModelInputSpec, ...]
+) -> tuple[torch.Tensor, ...]:
+    example_input = []
+    for spec in input_spec:
+        match spec.dim_order:
+            case torch.contiguous_format:
+                sample = torch.ones(spec.shape, dtype=spec.dtype)
+            case torch.channels_last:
+                sample = torch.ones(spec.shape, dtype=spec.dtype).to(
+                    memory_format=torch.channels_last
+                )
+            case _:
+                raise ValueError(f"Unsupported dim_order: {spec.dim_order}")
+        # noinspection PyUnboundLocalVariable
+        example_input.append(sample)
+
+    return tuple(example_input)
 
 
 def to_quantized_edge_program(
     model: torch.nn.Module,
-    input_spec: tuple[ModelInputSpec, ...] | tuple[int, ...] | list[tuple[int, ...]],
+    input_spec: list[ModelInputSpec] | tuple[int, ...] | list[tuple[int, ...]],
     operators_not_to_delegate: list[str] = None,
-    get_calibration_inputs_fn: Callable[
-        [tuple[ModelInputSpec, ...]], list[tuple[torch.Tensor, ...]]
-    ] = get_random_calibration_inputs,
+    get_calibration_inputs_fn: GetCalibrationInputsFn = get_random_calibration_inputs,
     target: str = "imxrt700",
     use_qat: bool = False,
+    train_fn: Callable[[torch.fx.GraphModule], None] | None = None,
     remove_quant_io_ops: bool = False,
     custom_delegation_options: CustomDelegationOptions = CustomDelegationOptions(),  # noqa B008
     get_quantizer_fn: Callable[[], Quantizer] | None = None,
@@ -131,15 +189,16 @@ def to_quantized_edge_program(
     fetch_constants_to_sram: bool = False,
     dump_kernel_selection_code: bool = False,
     use_new_flow_neutron_c: bool = False,
+    delegate_to_npu=True,
 ) -> EdgeProgramManager:
     _neutron_target_spec = NeutronTargetSpec(target)
     if get_quantizer_fn is None:
         get_quantizer_fn = partial(
             _get_default_quantizer, _neutron_target_spec, use_qat
         )
-
-    calibration_inputs = get_calibration_inputs_fn(to_model_input_spec(input_spec))
-    example_input = calibration_inputs[0]
+    input_spec = to_model_input_spec(input_spec)
+    calibration_inputs = get_calibration_inputs_fn(input_spec)
+    example_input = _get_example_input(input_spec)
 
     # Make sure the model is in the evaluation mode.
     model.eval()
@@ -151,6 +210,7 @@ def to_quantized_edge_program(
         calibration_inputs=calibration_inputs,
         quantizer=get_quantizer_fn(),
         is_qat=use_qat,
+        train_fn=train_fn,
     )
 
     # List of operators to not decompose during the lowering.
@@ -166,15 +226,18 @@ def to_quantized_edge_program(
     post_quant_state_dict = (
         exir_program_aten__module_quant.state_dict() if use_quant_state_dict else None
     )
-    partitioners = [
-        NeutronPartitioner(
-            compile_spec,
-            _neutron_target_spec,
-            custom_delegation_options,
-            post_quant_state_dict,
-            preserve_ops=preserve_ops,
-        )
-    ]
+    if delegate_to_npu:
+        partitioners = [
+            NeutronPartitioner(
+                compile_spec,
+                _neutron_target_spec,
+                custom_delegation_options,
+                post_quant_state_dict,
+                preserve_ops=preserve_ops,
+            )
+        ]
+    else:
+        partitioners = []
 
     edge_program_manager = to_edge_transform_and_lower(
         export(exir_program_aten__module_quant, example_input, strict=True),
@@ -205,13 +268,31 @@ def to_quantized_executorch_program(
     model: torch.nn.Module,
     input_spec: tuple[ModelInputSpec, ...] | tuple[int, ...] | list[tuple[int, ...]],
     use_qat: bool = False,
+    train_fn: Callable[[torch.fx.GraphModule], None] | None = None,
     use_neutron_for_format_conversion: bool = True,
+    dataset_dir: str | None = None,
+    delegate_to_npu=True,
+    use_new_flow_neutron_c: bool = False,
 ) -> ExecutorchProgramManager:
+    if dataset_dir:
+        # Extract calibration data from a directory.
+        get_calibration_inputs_fn = {
+            "get_calibration_inputs_fn": get_calibration_inputs_fn_from_dataset_dir(
+                dataset_dir
+            )
+        }
+    else:
+        get_calibration_inputs_fn = {}  # Use default parameter value.
+
     edge_program_manager = to_quantized_edge_program(
         model,
         input_spec,
         use_qat=use_qat,
+        train_fn=train_fn,
         use_neutron_for_format_conversion=use_neutron_for_format_conversion,
+        delegate_to_npu=delegate_to_npu,
+        use_new_flow_neutron_c=use_new_flow_neutron_c,
+        **get_calibration_inputs_fn,
     )
 
     return edge_program_manager.to_executorch(
diff --git a/backends/nxp/tests_models/dataset_creator.py b/backends/nxp/tests_models/dataset_creator.py
index 361f3a0889d..e0fc3444dbf 100644
--- a/backends/nxp/tests_models/dataset_creator.py
+++ b/backends/nxp/tests_models/dataset_creator.py
@@ -14,8 +14,8 @@
 import numpy as np
 import torch
 from executorch.backends.nxp.backend.ir.converter.conversion import translator
+from executorch.backends.nxp.tests.executorch_pipeline import ModelInputSpec
 from executorch.backends.nxp.tests_models.calibration_dataset import CalibrationDataset
-from executorch.backends.nxp.tests_models.model_input_spec import ModelInputSpec
 from torch import Tensor
 
 
diff --git a/backends/nxp/tests_models/executors.py b/backends/nxp/tests_models/executors.py
index 5209a57fac5..450932c14d1 100644
--- a/backends/nxp/tests_models/executors.py
+++ b/backends/nxp/tests_models/executors.py
@@ -19,19 +19,20 @@
 from executorch.backends.nxp.backend.edge_helper import is_channels_last_dim_order
 from executorch.backends.nxp.backend.ir.converter.conversion import translator
 from executorch.backends.nxp.neutron_partitioner import NeutronPartitioner
+from executorch.backends.nxp.tests.executorch_pipeline import (
+    get_calibration_inputs_fn_from_dataset_dir,
+    ModelInputSpec,
+    to_quantized_edge_program,
+    to_quantized_executorch_program,
+)
 from executorch.backends.nxp.tests_models.config_importer import test_config
 from executorch.backends.nxp.tests_models.dataset_creator import RandomDatasetCreator
 from executorch.backends.nxp.tests_models.graph_verifier import GraphVerifier
-from executorch.backends.nxp.tests_models.model_input_spec import ModelInputSpec
 from executorch.backends.nxp.tests_models.model_output_comparator import (
     AllCloseOutputComparator,
 )
 from executorch.backends.nxp.tests_models.outputs_dir_importer import outputs_dir
-from executorch.backends.nxp.tests_models.utils import (
-    save_pte_program,
-    to_quantized_edge_program,
-    to_quantized_executorch_program,
-)
+from executorch.backends.nxp.tests_models.utils import save_pte_program
 from executorch.devtools.visualization.visualization_utils import (
     visualize_with_clusters,
 )
@@ -113,7 +114,7 @@ def wrapper(*args, **kwargs):
         delegated_program = to_quantized_executorch_program(
             model,
             input_spec,
-            calibration_dataset_dir,
+            dataset_dir=calibration_dataset_dir,
             delegate_to_npu=True,
             use_qat=use_qat,
             train_fn=train_fn,
@@ -175,7 +176,7 @@ def _run_non_delegated_executorch_program(
     non_delegated_program = to_quantized_executorch_program(
         model,
         input_spec,
-        calibration_dataset_dir,
+        dataset_dir=calibration_dataset_dir,
         delegate_to_npu=False,
         use_qat=use_qat,
         train_fn=train_fn,
@@ -463,7 +464,9 @@ def convert_run_compare(
                 to_quantized_edge_program(
                     model_to_not_delegate,
                     input_spec,
-                    calibration_dataset_dir,
+                    get_calibration_inputs_fn=get_calibration_inputs_fn_from_dataset_dir(
+                        calibration_dataset_dir
+                    ),
                     delegate_to_npu=False,
                     use_qat=use_qat,
                     train_fn=train_fn,
diff --git a/backends/nxp/tests_models/model_input_spec.py b/backends/nxp/tests_models/model_input_spec.py
deleted file mode 100644
index c96ca53b8da..00000000000
--- a/backends/nxp/tests_models/model_input_spec.py
+++ /dev/null
@@ -1,18 +0,0 @@
-# Copyright 2025-2026 NXP
-#
-# This source code is licensed under the BSD-style license found in the
-# LICENSE file in the root directory of this source tree.
-
-from dataclasses import dataclass
-
-import numpy as np
-import torch
-from torch import memory_format
-
-
-@dataclass
-class ModelInputSpec:
-    shape: tuple[int, ...]
-    type: np.dtype = np.float32
-    dtype: torch.dtype = torch.float32
-    dim_order: memory_format = torch.contiguous_format
diff --git a/backends/nxp/tests_models/test_cifarnet.py b/backends/nxp/tests_models/test_cifarnet.py
index 97e7125f2fe..2890b92ae4a 100644
--- a/backends/nxp/tests_models/test_cifarnet.py
+++ b/backends/nxp/tests_models/test_cifarnet.py
@@ -7,6 +7,7 @@
 
 import pytest
 import torch
+from executorch.backends.nxp.tests.executorch_pipeline import ModelInputSpec
 
 from executorch.backends.nxp.tests_models.config_importer import test_config
 from executorch.backends.nxp.tests_models.dataset_creator import CopyDatasetCreator
@@ -19,7 +20,6 @@
     BaseGraphVerifier,
     NonDelegatedNode,
 )
-from executorch.backends.nxp.tests_models.model_input_spec import ModelInputSpec
 from executorch.backends.nxp.tests_models.model_output_comparator import (
     NumericalStatsOutputComparator,
 )
diff --git a/backends/nxp/tests_models/utils.py b/backends/nxp/tests_models/utils.py
index 06e09827746..c210d9db8bc 100644
--- a/backends/nxp/tests_models/utils.py
+++ b/backends/nxp/tests_models/utils.py
@@ -7,222 +7,11 @@
 
 import logging
 import os
-from typing import Any, Callable, Dict, Optional, Tuple, Union
 
-import executorch.exir as exir
-
-import numpy as np
-import torch
-
-from executorch.backends.nxp.aten_passes.fuse_batch_norm_with_linear_pass import (
-    FuseBatchNormWithLinearPass,
-)
-from executorch.backends.nxp.aten_passes.simulated_linear_bn_fusion_passes import (
-    AddSimulatedLinearBatchNormFusionQATPass,
-    RemoveSimulatedLinearBatchNormFusionQATPass,
-)
-from executorch.backends.nxp.backend.neutron_target_spec import NeutronTargetSpec
-from executorch.backends.nxp.edge_passes.neutron_edge_pass_manager import (
-    NeutronEdgePassManager,
-)
-from executorch.backends.nxp.neutron_partitioner import NeutronPartitioner
-
-from executorch.backends.nxp.nxp_backend import (
-    core_aten_ops_exception_list,
-    generate_neutron_compile_spec,
-)
-from executorch.backends.nxp.quantizer.neutron_quantizer import NeutronQuantizer
-from executorch.backends.nxp.tests_models.model_input_spec import ModelInputSpec
 from executorch.devtools.visualization.visualization_utils import (
     visualize_with_clusters,
 )
-from executorch.exir import (
-    EdgeProgramManager,
-    ExecutorchProgramManager,
-    to_edge_transform_and_lower,
-)
-from executorch.exir.capture._config import EdgeCompileConfig, ExecutorchBackendConfig
-from executorch.exir.tracer import Value
-from torch.export import export, ExportedProgram
-from torchao.quantization.pt2e import (
-    move_exported_model_to_eval,
-    move_exported_model_to_train,
-)
-from torchao.quantization.pt2e.quantize_pt2e import (
-    convert_pt2e,
-    prepare_pt2e,
-    prepare_qat_pt2e,
-)
-
-_EDGE_COMPILE_CONFIG = exir.EdgeCompileConfig(
-    _check_ir_validity=True,
-    _skip_dim_order=True,  # TODO(T189114319): Reuse dim order op after solving the ios oss issue
-)
-
-
-def to_quantized_edge_program(
-    model: torch.nn.Module,
-    input_spec: list[ModelInputSpec],
-    dataset_dir,
-    delegate_to_npu=True,
-    use_qat: bool = False,
-    train_fn: Callable[[torch.fx.GraphModule], None] | None = None,
-    use_new_flow_neutron_c: bool = False,
-) -> EdgeProgramManager:
-    assert isinstance(input_spec, list) and all(
-        isinstance(spec, ModelInputSpec) for spec in input_spec
-    ), "Input_spec must be a list of ModelInputSpec."
-
-    example_input = []
-    for spec in input_spec:
-        match spec.dim_order:
-            case torch.contiguous_format:
-                sample = torch.ones(spec.shape, dtype=spec.dtype)
-            case torch.channels_last:
-                sample = torch.ones(spec.shape, dtype=spec.dtype).to(
-                    memory_format=torch.channels_last
-                )
-            case _:
-                raise ValueError(f"Unsupported dim_order: {spec.dim_order}")
-        # noinspection PyUnboundLocalVariable
-        example_input.append(sample)
-
-    example_input = tuple(example_input)
-
-    exir_program_aten = torch.export.export(model, example_input, strict=True)
-    module = exir_program_aten.module()
-
-    neutron_target_spec = NeutronTargetSpec(target="imxrt700")
-
-    # Quantize model
-    quantizer = NeutronQuantizer(
-        neutron_target_spec=neutron_target_spec, is_qat=use_qat
-    )
-    if use_qat:
-        m = prepare_qat_pt2e(module, quantizer)
-        m = AddSimulatedLinearBatchNormFusionQATPass()(m).graph_module
-
-        if train_fn:
-            m = move_exported_model_to_train(m)
-            train_fn(m)
-
-        m = move_exported_model_to_eval(m)
-        m = RemoveSimulatedLinearBatchNormFusionQATPass()(m).graph_module
-        m = FuseBatchNormWithLinearPass()(m).graph_module
-    else:
-        m = prepare_pt2e(module, quantizer)
-
-    data = sorted(os.listdir(dataset_dir))
-    inputs_needed = len(input_spec)
-
-    # If the model is single-input, the following directory structure is used:
-    #   dataset_dir/data.bin (data.bin is *the* input)
-    # Else, if multi-input, the following directory structure is used:
-    #   dataset_dir/data/{.+}.bin (each .bin file is an input)
-
-    if not use_qat or (use_qat and not train_fn):
-        input_data = []
-        for path in data:
-            path = os.path.join(dataset_dir, path)
-            files = []
-
-            if os.path.isdir(path):
-                files = [os.path.join(path, x) for x in sorted(os.listdir(path))]
-            else:
-                files.append(path)
-
-            for idx, file in enumerate(files):
-                if len(input_data) == inputs_needed:
-                    break
-
-                tensor = np.fromfile(file, dtype=input_spec[idx].type).reshape(
-                    input_spec[idx].shape
-                )
-                input_data += (torch.from_numpy(tensor),)
-                continue
-
-            if len(input_data) < inputs_needed:
-                continue
-
-            m(*input_data)
-            input_data.clear()
-
-    exir_program_aten_quant = convert_pt2e(m)
-
-    # To ATen
-    core_aten_ep = _to_core_aten(
-        exir_program_aten_quant, example_input, None, verbose=True
-    )
-
-    partitioners = (
-        (
-            [
-                NeutronPartitioner(
-                    generate_neutron_compile_spec(
-                        "imxrt700", use_new_flow_neutron_c=use_new_flow_neutron_c
-                    ),
-                    neutron_target_spec=neutron_target_spec,
-                    post_quantization_state_dict=exir_program_aten_quant.state_dict(),
-                )
-            ]
-        )
-        if delegate_to_npu
-        else []
-    )
-
-    edge_program_manager = to_edge_transform_and_lower(
-        core_aten_ep,
-        transform_passes=NeutronEdgePassManager(),
-        partitioner=partitioners,
-        compile_config=EdgeCompileConfig(
-            _core_aten_ops_exception_list=core_aten_ops_exception_list
-        ),
-    )
-
-    return edge_program_manager
-
-
-def to_quantized_executorch_program(
-    model: torch.nn.Module,
-    input_spec,
-    dataset_dir: str,
-    delegate_to_npu=True,
-    use_qat: bool = False,
-    train_fn: Callable[[torch.fx.GraphModule], None] | None = None,
-    use_new_flow_neutron_c: bool = False,
-) -> ExecutorchProgramManager:
-    edge_program_manager = to_quantized_edge_program(
-        model,
-        input_spec,
-        dataset_dir,
-        delegate_to_npu,
-        use_qat=use_qat,
-        train_fn=train_fn,
-        use_new_flow_neutron_c=use_new_flow_neutron_c,
-    )
-
-    return edge_program_manager.to_executorch(
-        config=ExecutorchBackendConfig(extract_delegate_segments=False)
-    )
-
-
-def _to_core_aten(
-    model: Union[torch.fx.GraphModule, torch.nn.Module],
-    example_inputs: Tuple[Value, ...],
-    dynamic_shapes: Optional[Union[Dict[str, Any], Tuple[Any]]] = None,
-    verbose=True,
-) -> ExportedProgram:
-    # post autograd export. eventually this will become .to_core_aten
-    if not isinstance(model, torch.fx.GraphModule) and not isinstance(
-        model, torch.nn.Module
-    ):
-        raise ValueError(
-            f"Expected passed in model to be an instance of fx.GraphModule, got {type(model)}"
-        )
-    core_aten_ep = export(model, example_inputs, dynamic_shapes=dynamic_shapes)
-    if verbose:
-        logging.info(f"Core ATen graph:\n{core_aten_ep.graph}")
-    return core_aten_ep
+from executorch.exir import ExecutorchProgramManager
 
 
 def save_pte_program(

From eef4f6d606991774456dd244ee885adcf75afe38 Mon Sep 17 00:00:00 2001
From: Martin Pavella <martin.pavella@nxp.com>
Date: Fri, 24 Apr 2026 12:36:39 +0200
Subject: [PATCH 2/4] Move generic tests into a separate subdirectory.

---
 __init__.py                                                   | 0
 backends/nxp/__init__.py                                      | 0
 backends/nxp/tests/__init__.py                                | 0
 backends/nxp/tests/generic_tests/__init__.py                  | 0
 backends/nxp/tests/{ => generic_tests}/test_aot_example.py    | 4 ++--
 .../nxp/tests/{ => generic_tests}/test_batch_norm_fusion.py   | 0
 .../{ => generic_tests}/test_context_sensitive_delegation.py  | 0
 .../nxp/tests/{ => generic_tests}/test_convert_div_to_mul.py  | 0
 .../{ => generic_tests}/test_decompose_split_to_slices.py     | 0
 backends/nxp/tests/{ => generic_tests}/test_gru_splitting.py  | 0
 backends/nxp/tests/{ => generic_tests}/test_integration.py    | 0
 .../nxp/tests/{ => generic_tests}/test_kernel_selection.py    | 0
 .../tests/{ => generic_tests}/test_linear_and_add_fusion.py   | 0
 .../test_move_activation_before_concatenation.py              | 0
 .../nxp/tests/{ => generic_tests}/test_neutron_backend.py     | 0
 .../{ => generic_tests}/test_neutron_backend_executor.py      | 0
 .../{ => generic_tests}/test_neutron_converter_manager.py     | 0
 .../tests/{ => generic_tests}/test_node_format_inference.py   | 0
 .../nxp/tests/{ => generic_tests}/test_operator_selector.py   | 0
 .../tests/{ => generic_tests}/test_per_channel_conversion.py  | 0
 .../nxp/tests/{ => generic_tests}/test_qdq_clustering_conv.py | 0
 backends/nxp/tests/{ => generic_tests}/test_quantizer.py      | 0
 .../nxp/tests/{ => generic_tests}/test_removing_dead_code.py  | 0
 .../test_removing_nodes_with_known_outputs.py                 | 0
 .../tests/{ => generic_tests}/test_split_group_convolution.py | 0
 25 files changed, 2 insertions(+), 2 deletions(-)
 create mode 100644 __init__.py
 create mode 100644 backends/nxp/__init__.py
 create mode 100644 backends/nxp/tests/__init__.py
 create mode 100644 backends/nxp/tests/generic_tests/__init__.py
 rename backends/nxp/tests/{ => generic_tests}/test_aot_example.py (97%)
 rename backends/nxp/tests/{ => generic_tests}/test_batch_norm_fusion.py (100%)
 rename backends/nxp/tests/{ => generic_tests}/test_context_sensitive_delegation.py (100%)
 rename backends/nxp/tests/{ => generic_tests}/test_convert_div_to_mul.py (100%)
 rename backends/nxp/tests/{ => generic_tests}/test_decompose_split_to_slices.py (100%)
 rename backends/nxp/tests/{ => generic_tests}/test_gru_splitting.py (100%)
 rename backends/nxp/tests/{ => generic_tests}/test_integration.py (100%)
 rename backends/nxp/tests/{ => generic_tests}/test_kernel_selection.py (100%)
 rename backends/nxp/tests/{ => generic_tests}/test_linear_and_add_fusion.py (100%)
 rename backends/nxp/tests/{ => generic_tests}/test_move_activation_before_concatenation.py (100%)
 rename backends/nxp/tests/{ => generic_tests}/test_neutron_backend.py (100%)
 rename backends/nxp/tests/{ => generic_tests}/test_neutron_backend_executor.py (100%)
 rename backends/nxp/tests/{ => generic_tests}/test_neutron_converter_manager.py (100%)
 rename backends/nxp/tests/{ => generic_tests}/test_node_format_inference.py (100%)
 rename backends/nxp/tests/{ => generic_tests}/test_operator_selector.py (100%)
 rename backends/nxp/tests/{ => generic_tests}/test_per_channel_conversion.py (100%)
 rename backends/nxp/tests/{ => generic_tests}/test_qdq_clustering_conv.py (100%)
 rename backends/nxp/tests/{ => generic_tests}/test_quantizer.py (100%)
 rename backends/nxp/tests/{ => generic_tests}/test_removing_dead_code.py (100%)
 rename backends/nxp/tests/{ => generic_tests}/test_removing_nodes_with_known_outputs.py (100%)
 rename backends/nxp/tests/{ => generic_tests}/test_split_group_convolution.py (100%)

diff --git a/__init__.py b/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/backends/nxp/__init__.py b/backends/nxp/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/backends/nxp/tests/__init__.py b/backends/nxp/tests/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/backends/nxp/tests/generic_tests/__init__.py b/backends/nxp/tests/generic_tests/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/backends/nxp/tests/test_aot_example.py b/backends/nxp/tests/generic_tests/test_aot_example.py
similarity index 97%
rename from backends/nxp/tests/test_aot_example.py
rename to backends/nxp/tests/generic_tests/test_aot_example.py
index e2864f2c2c3..893041fe372 100644
--- a/backends/nxp/tests/test_aot_example.py
+++ b/backends/nxp/tests/generic_tests/test_aot_example.py
@@ -15,8 +15,8 @@
 def test_aot_example__mobilenet_v2():
     """Test that mobilenet can be lowered to Neutron backend via `aot_neutron_compile.py` and all ops are delegated."""
 
-    # Find the executorch root directory (4 levels up from this test file)
-    executorch_root = Path(__file__).parent.parent.parent.parent
+    # Find the executorch root directory (5 levels up from this test file)
+    executorch_root = Path(__file__).parent.parent.parent.parent.parent
     assert executorch_root.exists(), f"Executorch root not found at {executorch_root}"
 
     # Run the compilation script as a module (like run_aot_example.sh does)
diff --git a/backends/nxp/tests/test_batch_norm_fusion.py b/backends/nxp/tests/generic_tests/test_batch_norm_fusion.py
similarity index 100%
rename from backends/nxp/tests/test_batch_norm_fusion.py
rename to backends/nxp/tests/generic_tests/test_batch_norm_fusion.py
diff --git a/backends/nxp/tests/test_context_sensitive_delegation.py b/backends/nxp/tests/generic_tests/test_context_sensitive_delegation.py
similarity index 100%
rename from backends/nxp/tests/test_context_sensitive_delegation.py
rename to backends/nxp/tests/generic_tests/test_context_sensitive_delegation.py
diff --git a/backends/nxp/tests/test_convert_div_to_mul.py b/backends/nxp/tests/generic_tests/test_convert_div_to_mul.py
similarity index 100%
rename from backends/nxp/tests/test_convert_div_to_mul.py
rename to backends/nxp/tests/generic_tests/test_convert_div_to_mul.py
diff --git a/backends/nxp/tests/test_decompose_split_to_slices.py b/backends/nxp/tests/generic_tests/test_decompose_split_to_slices.py
similarity index 100%
rename from backends/nxp/tests/test_decompose_split_to_slices.py
rename to backends/nxp/tests/generic_tests/test_decompose_split_to_slices.py
diff --git a/backends/nxp/tests/test_gru_splitting.py b/backends/nxp/tests/generic_tests/test_gru_splitting.py
similarity index 100%
rename from backends/nxp/tests/test_gru_splitting.py
rename to backends/nxp/tests/generic_tests/test_gru_splitting.py
diff --git a/backends/nxp/tests/test_integration.py b/backends/nxp/tests/generic_tests/test_integration.py
similarity index 100%
rename from backends/nxp/tests/test_integration.py
rename to backends/nxp/tests/generic_tests/test_integration.py
diff --git a/backends/nxp/tests/test_kernel_selection.py b/backends/nxp/tests/generic_tests/test_kernel_selection.py
similarity index 100%
rename from backends/nxp/tests/test_kernel_selection.py
rename to backends/nxp/tests/generic_tests/test_kernel_selection.py
diff --git a/backends/nxp/tests/test_linear_and_add_fusion.py b/backends/nxp/tests/generic_tests/test_linear_and_add_fusion.py
similarity index 100%
rename from backends/nxp/tests/test_linear_and_add_fusion.py
rename to backends/nxp/tests/generic_tests/test_linear_and_add_fusion.py
diff --git a/backends/nxp/tests/test_move_activation_before_concatenation.py b/backends/nxp/tests/generic_tests/test_move_activation_before_concatenation.py
similarity index 100%
rename from backends/nxp/tests/test_move_activation_before_concatenation.py
rename to backends/nxp/tests/generic_tests/test_move_activation_before_concatenation.py
diff --git a/backends/nxp/tests/test_neutron_backend.py b/backends/nxp/tests/generic_tests/test_neutron_backend.py
similarity index 100%
rename from backends/nxp/tests/test_neutron_backend.py
rename to backends/nxp/tests/generic_tests/test_neutron_backend.py
diff --git a/backends/nxp/tests/test_neutron_backend_executor.py b/backends/nxp/tests/generic_tests/test_neutron_backend_executor.py
similarity index 100%
rename from backends/nxp/tests/test_neutron_backend_executor.py
rename to backends/nxp/tests/generic_tests/test_neutron_backend_executor.py
diff --git a/backends/nxp/tests/test_neutron_converter_manager.py b/backends/nxp/tests/generic_tests/test_neutron_converter_manager.py
similarity index 100%
rename from backends/nxp/tests/test_neutron_converter_manager.py
rename to backends/nxp/tests/generic_tests/test_neutron_converter_manager.py
diff --git a/backends/nxp/tests/test_node_format_inference.py b/backends/nxp/tests/generic_tests/test_node_format_inference.py
similarity index 100%
rename from backends/nxp/tests/test_node_format_inference.py
rename to backends/nxp/tests/generic_tests/test_node_format_inference.py
diff --git a/backends/nxp/tests/test_operator_selector.py b/backends/nxp/tests/generic_tests/test_operator_selector.py
similarity index 100%
rename from backends/nxp/tests/test_operator_selector.py
rename to backends/nxp/tests/generic_tests/test_operator_selector.py
diff --git a/backends/nxp/tests/test_per_channel_conversion.py b/backends/nxp/tests/generic_tests/test_per_channel_conversion.py
similarity index 100%
rename from backends/nxp/tests/test_per_channel_conversion.py
rename to backends/nxp/tests/generic_tests/test_per_channel_conversion.py
diff --git a/backends/nxp/tests/test_qdq_clustering_conv.py b/backends/nxp/tests/generic_tests/test_qdq_clustering_conv.py
similarity index 100%
rename from backends/nxp/tests/test_qdq_clustering_conv.py
rename to backends/nxp/tests/generic_tests/test_qdq_clustering_conv.py
diff --git a/backends/nxp/tests/test_quantizer.py b/backends/nxp/tests/generic_tests/test_quantizer.py
similarity index 100%
rename from backends/nxp/tests/test_quantizer.py
rename to backends/nxp/tests/generic_tests/test_quantizer.py
diff --git a/backends/nxp/tests/test_removing_dead_code.py b/backends/nxp/tests/generic_tests/test_removing_dead_code.py
similarity index 100%
rename from backends/nxp/tests/test_removing_dead_code.py
rename to backends/nxp/tests/generic_tests/test_removing_dead_code.py
diff --git a/backends/nxp/tests/test_removing_nodes_with_known_outputs.py b/backends/nxp/tests/generic_tests/test_removing_nodes_with_known_outputs.py
similarity index 100%
rename from backends/nxp/tests/test_removing_nodes_with_known_outputs.py
rename to backends/nxp/tests/generic_tests/test_removing_nodes_with_known_outputs.py
diff --git a/backends/nxp/tests/test_split_group_convolution.py b/backends/nxp/tests/generic_tests/test_split_group_convolution.py
similarity index 100%
rename from backends/nxp/tests/test_split_group_convolution.py
rename to backends/nxp/tests/generic_tests/test_split_group_convolution.py

From c213c9f64b2b04d4431982f79ead2e2780b7682d Mon Sep 17 00:00:00 2001
From: Martin Pavella <martin.pavella@nxp.com>
Date: Fri, 24 Apr 2026 15:39:13 +0200
Subject: [PATCH 3/4] Move files from `/tests_models` to `/tests`.

---
 backends/nxp/quantizer/utils.py               |  2 +-
 .../calibration_dataset.py                    |  0
 .../nxp/{tests_models => tests}/config.py     |  2 +-
 .../config_importer.py                        |  2 +-
 .../nxp/{tests_models => tests}/conftest.py   |  2 +-
 .../dataset_creator.py                        |  2 +-
 backends/nxp/tests/executorch_pipeline.py     |  8 +++----
 .../generic_tests}/test_cifarnet.py           | 21 ++++++++-----------
 .../{tests_models => tests}/graph_verifier.py |  0
 .../model_output_comparator.py                |  0
 .../neutron-imxrt700.ini                      |  0
 .../executors.py => tests/nsys_testing.py}    | 18 ++++++++--------
 .../{tests_models => tests}/outputs_dir.py    |  0
 .../outputs_dir_importer.py                   |  2 +-
 backends/nxp/{tests_models => tests}/utils.py |  0
 backends/nxp/tests_models/__init__.py         |  4 ----
 16 files changed, 28 insertions(+), 35 deletions(-)
 rename backends/nxp/{tests_models => tests}/calibration_dataset.py (100%)
 rename backends/nxp/{tests_models => tests}/config.py (94%)
 rename backends/nxp/{tests_models => tests}/config_importer.py (82%)
 rename backends/nxp/{tests_models => tests}/conftest.py (92%)
 rename backends/nxp/{tests_models => tests}/dataset_creator.py (99%)
 rename backends/nxp/{tests_models => tests/generic_tests}/test_cifarnet.py (86%)
 rename backends/nxp/{tests_models => tests}/graph_verifier.py (100%)
 rename backends/nxp/{tests_models => tests}/model_output_comparator.py (100%)
 rename backends/nxp/{tests_models => tests}/neutron-imxrt700.ini (100%)
 rename backends/nxp/{tests_models/executors.py => tests/nsys_testing.py} (97%)
 rename backends/nxp/{tests_models => tests}/outputs_dir.py (100%)
 rename backends/nxp/{tests_models => tests}/outputs_dir_importer.py (82%)
 rename backends/nxp/{tests_models => tests}/utils.py (100%)
 delete mode 100644 backends/nxp/tests_models/__init__.py

diff --git a/backends/nxp/quantizer/utils.py b/backends/nxp/quantizer/utils.py
index 09ee0e5b2f8..cd403868a96 100644
--- a/backends/nxp/quantizer/utils.py
+++ b/backends/nxp/quantizer/utils.py
@@ -1,5 +1,5 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
-# Copyright 2024-2025 NXP
+# Copyright 2024-2026 NXP
 # All rights reserved.
 #
 # This source code is licensed under the BSD-style license found in the
diff --git a/backends/nxp/tests_models/calibration_dataset.py b/backends/nxp/tests/calibration_dataset.py
similarity index 100%
rename from backends/nxp/tests_models/calibration_dataset.py
rename to backends/nxp/tests/calibration_dataset.py
diff --git a/backends/nxp/tests_models/config.py b/backends/nxp/tests/config.py
similarity index 94%
rename from backends/nxp/tests_models/config.py
rename to backends/nxp/tests/config.py
index 7a5f108c9b7..a9f77c517e3 100644
--- a/backends/nxp/tests_models/config.py
+++ b/backends/nxp/tests/config.py
@@ -21,7 +21,7 @@
 
 NSYS_PATH = pathlib.Path(shutil.which("nsys"))
 NSYS_CONFIG_PATH = os.path.join(
-    PROJECT_DIR, "backends", "nxp", "tests_models", "neutron-imxrt700.ini"
+    PROJECT_DIR, "backends", "nxp", "tests", "neutron-imxrt700.ini"
 )
 NSYS_FIRMWARE_PATH = os.path.join(
     os.path.dirname(eiq_neutron_sdk.__file__),
diff --git a/backends/nxp/tests_models/config_importer.py b/backends/nxp/tests/config_importer.py
similarity index 82%
rename from backends/nxp/tests_models/config_importer.py
rename to backends/nxp/tests/config_importer.py
index 751ed9bd751..286e4fd0f73 100644
--- a/backends/nxp/tests_models/config_importer.py
+++ b/backends/nxp/tests/config_importer.py
@@ -12,6 +12,6 @@
 
     logger.debug("Importing from executorch-integration")
 except ImportError:
-    import executorch.backends.nxp.tests_models.config as test_config  # noqa F401
+    import executorch.backends.nxp.tests.config as test_config  # noqa F401
 
     logger.debug("Importing from executorch")
diff --git a/backends/nxp/tests_models/conftest.py b/backends/nxp/tests/conftest.py
similarity index 92%
rename from backends/nxp/tests_models/conftest.py
rename to backends/nxp/tests/conftest.py
index f4c7a36092a..34fe343ca6a 100644
--- a/backends/nxp/tests_models/conftest.py
+++ b/backends/nxp/tests/conftest.py
@@ -8,7 +8,7 @@
 import pathlib
 import shutil
 
-from executorch.backends.nxp.tests_models.outputs_dir_importer import outputs_dir
+from executorch.backends.nxp.tests.outputs_dir_importer import outputs_dir
 
 
 def pytest_addoption(parser):
diff --git a/backends/nxp/tests_models/dataset_creator.py b/backends/nxp/tests/dataset_creator.py
similarity index 99%
rename from backends/nxp/tests_models/dataset_creator.py
rename to backends/nxp/tests/dataset_creator.py
index e0fc3444dbf..928c392bf17 100644
--- a/backends/nxp/tests_models/dataset_creator.py
+++ b/backends/nxp/tests/dataset_creator.py
@@ -14,8 +14,8 @@
 import numpy as np
 import torch
 from executorch.backends.nxp.backend.ir.converter.conversion import translator
+from executorch.backends.nxp.tests.calibration_dataset import CalibrationDataset
 from executorch.backends.nxp.tests.executorch_pipeline import ModelInputSpec
-from executorch.backends.nxp.tests_models.calibration_dataset import CalibrationDataset
 from torch import Tensor
 
 
diff --git a/backends/nxp/tests/executorch_pipeline.py b/backends/nxp/tests/executorch_pipeline.py
index 55c8b1e2b5f..57935b41c18 100644
--- a/backends/nxp/tests/executorch_pipeline.py
+++ b/backends/nxp/tests/executorch_pipeline.py
@@ -108,7 +108,7 @@ def to_model_input_spec(
         case list() if all(
             isinstance(input_shape, tuple) for input_shape in input_spec
         ):
-            return tuple([ModelInputSpec(spec) for spec in input_spec])
+            return tuple(ModelInputSpec(spec) for spec in input_spec)
         case _:
             raise TypeError(f"Unsupported type {type(input_spec)}")
 
@@ -175,7 +175,7 @@ def _get_example_input(
 
 def to_quantized_edge_program(
     model: torch.nn.Module,
-    input_spec: list[ModelInputSpec] | tuple[int, ...] | list[tuple[int, ...]],
+    input_spec: Iterable[ModelInputSpec] | tuple[int, ...] | list[tuple[int, ...]],
     operators_not_to_delegate: list[str] = None,
     get_calibration_inputs_fn: GetCalibrationInputsFn = get_random_calibration_inputs,
     target: str = "imxrt700",
@@ -266,7 +266,7 @@ def to_quantized_edge_program(
 
 def to_quantized_executorch_program(
     model: torch.nn.Module,
-    input_spec: tuple[ModelInputSpec, ...] | tuple[int, ...] | list[tuple[int, ...]],
+    input_spec: Iterable[ModelInputSpec] | tuple[int, ...] | list[tuple[int, ...]],
     use_qat: bool = False,
     train_fn: Callable[[torch.fx.GraphModule], None] | None = None,
     use_neutron_for_format_conversion: bool = True,
@@ -302,7 +302,7 @@ def to_quantized_executorch_program(
 
 def to_edge_program(
     model: nn.Module,
-    input_spec: tuple[ModelInputSpec, ...] | tuple[int, ...] | list[tuple[int, ...]],
+    input_spec: Iterable[ModelInputSpec] | tuple[int, ...] | list[tuple[int, ...]],
 ) -> EdgeProgramManager:
     calibration_inputs = get_random_calibration_inputs(to_model_input_spec(input_spec))
 
diff --git a/backends/nxp/tests_models/test_cifarnet.py b/backends/nxp/tests/generic_tests/test_cifarnet.py
similarity index 86%
rename from backends/nxp/tests_models/test_cifarnet.py
rename to backends/nxp/tests/generic_tests/test_cifarnet.py
index 2890b92ae4a..1d795c938fe 100644
--- a/backends/nxp/tests_models/test_cifarnet.py
+++ b/backends/nxp/tests/generic_tests/test_cifarnet.py
@@ -7,22 +7,19 @@
 
 import pytest
 import torch
-from executorch.backends.nxp.tests.executorch_pipeline import ModelInputSpec
-
-from executorch.backends.nxp.tests_models.config_importer import test_config
-from executorch.backends.nxp.tests_models.dataset_creator import CopyDatasetCreator
 
-from executorch.backends.nxp.tests_models.executors import (
-    convert_run_compare,
-    ReferenceModel,
-)
-from executorch.backends.nxp.tests_models.graph_verifier import (
+from executorch.backends.nxp.tests.config_importer import test_config
+from executorch.backends.nxp.tests.dataset_creator import CopyDatasetCreator
+from executorch.backends.nxp.tests.executorch_pipeline import ModelInputSpec
+from executorch.backends.nxp.tests.graph_verifier import (
     BaseGraphVerifier,
     NonDelegatedNode,
 )
-from executorch.backends.nxp.tests_models.model_output_comparator import (
+from executorch.backends.nxp.tests.model_output_comparator import (
     NumericalStatsOutputComparator,
 )
+
+from executorch.backends.nxp.tests.nsys_testing import lower_run_compare, ReferenceModel
 from executorch.examples.nxp.experimental.cifar_net.cifar_net import (
     CifarNet,
     store_test_data,
@@ -64,7 +61,7 @@ def test_cifarnet(mocker, cifar_test_files, channels_last):
     comparator = NumericalStatsOutputComparator(
         max_mse_error=1.0e-3, is_classification_task=True
     )
-    convert_run_compare(
+    lower_run_compare(
         model,
         [input_spec],
         dataset_creator=CopyDatasetCreator(cifar_test_files),
@@ -94,7 +91,7 @@ def test_cifarnet_qat(mocker, cifar_test_files):
     comparator = NumericalStatsOutputComparator(
         max_mse_error=8e-2, is_classification_task=True
     )
-    convert_run_compare(
+    lower_run_compare(
         model,
         input_shape,
         dataset_creator=CopyDatasetCreator(cifar_test_files),
diff --git a/backends/nxp/tests_models/graph_verifier.py b/backends/nxp/tests/graph_verifier.py
similarity index 100%
rename from backends/nxp/tests_models/graph_verifier.py
rename to backends/nxp/tests/graph_verifier.py
diff --git a/backends/nxp/tests_models/model_output_comparator.py b/backends/nxp/tests/model_output_comparator.py
similarity index 100%
rename from backends/nxp/tests_models/model_output_comparator.py
rename to backends/nxp/tests/model_output_comparator.py
diff --git a/backends/nxp/tests_models/neutron-imxrt700.ini b/backends/nxp/tests/neutron-imxrt700.ini
similarity index 100%
rename from backends/nxp/tests_models/neutron-imxrt700.ini
rename to backends/nxp/tests/neutron-imxrt700.ini
diff --git a/backends/nxp/tests_models/executors.py b/backends/nxp/tests/nsys_testing.py
similarity index 97%
rename from backends/nxp/tests_models/executors.py
rename to backends/nxp/tests/nsys_testing.py
index 450932c14d1..ff4247c4729 100644
--- a/backends/nxp/tests_models/executors.py
+++ b/backends/nxp/tests/nsys_testing.py
@@ -19,20 +19,20 @@
 from executorch.backends.nxp.backend.edge_helper import is_channels_last_dim_order
 from executorch.backends.nxp.backend.ir.converter.conversion import translator
 from executorch.backends.nxp.neutron_partitioner import NeutronPartitioner
+from executorch.backends.nxp.tests.config_importer import test_config
+from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator
 from executorch.backends.nxp.tests.executorch_pipeline import (
     get_calibration_inputs_fn_from_dataset_dir,
     ModelInputSpec,
     to_quantized_edge_program,
     to_quantized_executorch_program,
 )
-from executorch.backends.nxp.tests_models.config_importer import test_config
-from executorch.backends.nxp.tests_models.dataset_creator import RandomDatasetCreator
-from executorch.backends.nxp.tests_models.graph_verifier import GraphVerifier
-from executorch.backends.nxp.tests_models.model_output_comparator import (
+from executorch.backends.nxp.tests.graph_verifier import GraphVerifier
+from executorch.backends.nxp.tests.model_output_comparator import (
     AllCloseOutputComparator,
 )
-from executorch.backends.nxp.tests_models.outputs_dir_importer import outputs_dir
-from executorch.backends.nxp.tests_models.utils import save_pte_program
+from executorch.backends.nxp.tests.outputs_dir_importer import outputs_dir
+from executorch.backends.nxp.tests.utils import save_pte_program
 from executorch.devtools.visualization.visualization_utils import (
     visualize_with_clusters,
 )
@@ -368,7 +368,7 @@ def assert_NSYS():
     assert os.path.exists(NSYS_FIRMWARE_PATH)
 
 
-def convert_run_compare(
+def lower_run_compare(
     model: torch.nn.Module,
     input_spec: list[ModelInputSpec] | tuple,
     dlg_model_verifier: GraphVerifier,
@@ -506,7 +506,7 @@ def convert_run_compare(
     )
 
 
-def convert_run_compare_ptq_qat(
+def lower_run_compare_ptq_qat(
     model: torch.nn.Module,
     input_spec: list[ModelInputSpec] | tuple,
     dlg_model_verifier: GraphVerifier,
@@ -593,7 +593,7 @@ def convert_run_compare_ptq_qat(
 
 
 def _get_caller_name():
-    test_function_names = ["convert_run_compare", "convert_run_compare_ptq_qat"]
+    test_function_names = ["lower_run_compare", "lower_run_compare_ptq_qat"]
     for idx, frame in enumerate(inspect.stack()):
         if frame.function in test_function_names:
             # Look one index above to get caller
diff --git a/backends/nxp/tests_models/outputs_dir.py b/backends/nxp/tests/outputs_dir.py
similarity index 100%
rename from backends/nxp/tests_models/outputs_dir.py
rename to backends/nxp/tests/outputs_dir.py
diff --git a/backends/nxp/tests_models/outputs_dir_importer.py b/backends/nxp/tests/outputs_dir_importer.py
similarity index 82%
rename from backends/nxp/tests_models/outputs_dir_importer.py
rename to backends/nxp/tests/outputs_dir_importer.py
index 2234dabbbeb..c018123c949 100644
--- a/backends/nxp/tests_models/outputs_dir_importer.py
+++ b/backends/nxp/tests/outputs_dir_importer.py
@@ -12,6 +12,6 @@
 
     logger.debug("Importing from executorch-integration")
 except ImportError:
-    import executorch.backends.nxp.tests_models.outputs_dir as outputs_dir  # noqa F401
+    import executorch.backends.nxp.tests.outputs_dir as outputs_dir  # noqa F401
 
     logger.debug("Importing from executorch")
diff --git a/backends/nxp/tests_models/utils.py b/backends/nxp/tests/utils.py
similarity index 100%
rename from backends/nxp/tests_models/utils.py
rename to backends/nxp/tests/utils.py
diff --git a/backends/nxp/tests_models/__init__.py b/backends/nxp/tests_models/__init__.py
deleted file mode 100644
index 55dc5fccf45..00000000000
--- a/backends/nxp/tests_models/__init__.py
+++ /dev/null
@@ -1,4 +0,0 @@
-# Copyright 2026 NXP
-#
-# This source code is licensed under the BSD-style license found in the
-# LICENSE file in the root directory of this source tree.

From 0394b4d889e3f47142e258f2c4abe6ae970926ae Mon Sep 17 00:00:00 2001
From: Martin Pavella <martin.pavella@nxp.com>
Date: Mon, 27 Apr 2026 08:29:49 +0200
Subject: [PATCH 4/4] NXP backend: Test `avg_pool2d` with new Neutron flow.

---
 .../nxp/backend/custom_delegation_options.py  |  2 +
 .../ops_converters/avg_pool_2d_converter.py   | 40 +++++++++-
 backends/nxp/nxp_backend.py                   |  6 ++
 .../test_avg_pool2d_converter.py              | 75 ++++++++++++++++++-
 backends/nxp/tests/models.py                  |  6 +-
 5 files changed, 121 insertions(+), 8 deletions(-)

diff --git a/backends/nxp/backend/custom_delegation_options.py b/backends/nxp/backend/custom_delegation_options.py
index 18eadc0bbbf..8016bb4f14c 100644
--- a/backends/nxp/backend/custom_delegation_options.py
+++ b/backends/nxp/backend/custom_delegation_options.py
@@ -22,3 +22,5 @@ class CustomDelegationOptions:
     #  not create any NeutronGraph that can be called. This is done by the partitioner itself, and is not handled by
     #  the individual node converters.
     allow_no_op_partitions: bool = False
+
+    use_new_flow_neutron_c: bool = False
diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/avg_pool_2d_converter.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/avg_pool_2d_converter.py
index 99ae0a30dbb..505e85baf88 100644
--- a/backends/nxp/backend/ir/converter/node_converters/ops_converters/avg_pool_2d_converter.py
+++ b/backends/nxp/backend/ir/converter/node_converters/ops_converters/avg_pool_2d_converter.py
@@ -1,4 +1,4 @@
-# Copyright 2025 NXP
+# Copyright 2025-2026 NXP
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
@@ -21,6 +21,8 @@
 from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options import (
     average_pool_2d_options,
 )
+
+from executorch.backends.nxp.backend.neutron_target_spec import NeutronTargetSpec
 from torch.fx import Node
 from torch.nn import Parameter
 
@@ -53,6 +55,27 @@ def _is_supported_in_IR(
 
         return True
 
+    @staticmethod
+    def _is_supported_on_target(
+        node: Node,
+        neutron_target_spec: NeutronTargetSpec,
+        parameters_mapping: dict[str, Parameter],
+        custom_delegation_options: CustomDelegationOptions,
+    ) -> bool:
+        kernel = node.args[1]
+        stride = node.args[2]
+
+        if custom_delegation_options.use_new_flow_neutron_c:
+            # Requirements specified by the new Neutron flow documentation.
+
+            if any(k > 4096 for k in kernel):
+                return False
+
+            if any(s > 4096 for s in stride):
+                return False
+
+        return True
+
     # noinspection PyMethodMayBeStatic
     def _convert_2d_avg_pool(
         self, kernel_size, stride, padding, t_op: tflite_model.Operator
@@ -85,10 +108,19 @@ def _convert_2d_avg_pool(
 
         return ops.flatten()
 
-    # AvgPool2d Node format: (Tensor self, int[2] kernel_size, int[2] stride=[], int[2] padding=0, bool ceil_mode=False
-    #                         bool count_include_pad=True, int? divisor_override=None)
     def convert(self, node: Node):
-        """Convert 'avg_pool2d' operator to TFLite 'AveragePool2D'."""
+        """Convert 'avg_pool2d' operator to TFLite 'AveragePool2D'.
+        The ExecuTorch schema is:
+            aten.avg_pool2d(
+                Tensor self,
+                int[2] kernel_size,
+                int[2] stride=[],
+                int[2] padding=0,
+                bool ceil_mode=False
+                bool count_include_pad=True,
+                int? divisor_override=None
+            )
+        """
         self.assert_convertible(node)
 
         kernel_size = node.args[1]
diff --git a/backends/nxp/nxp_backend.py b/backends/nxp/nxp_backend.py
index f7ecb8a908e..f5e89823ee2 100644
--- a/backends/nxp/nxp_backend.py
+++ b/backends/nxp/nxp_backend.py
@@ -15,6 +15,9 @@
 import numpy as np
 import torch
 
+from executorch.backends.nxp.backend.custom_delegation_options import (
+    CustomDelegationOptions,
+)
 from executorch.backends.nxp.backend.data_format import DataFormat
 from executorch.backends.nxp.backend.edge_program_converter import (
     EdgeProgramToIRConverter,
@@ -229,6 +232,9 @@ def preprocess(  # noqa C901
                 edge_program,
                 neutron_target_spec=NeutronTargetSpec(target),
                 conversion_config=conversion_config,
+                custom_delegation_options=CustomDelegationOptions(
+                    use_new_flow_neutron_c=use_new_flow_neutron_c
+                ),
             )
 
             neutron_model = NeutronConverterManager(dump_kernel_selection_code).convert(
diff --git a/backends/nxp/tests/ir/converter/node_converter/test_avg_pool2d_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_avg_pool2d_converter.py
index 73d1a060af1..f955cd075e1 100644
--- a/backends/nxp/tests/ir/converter/node_converter/test_avg_pool2d_converter.py
+++ b/backends/nxp/tests/ir/converter/node_converter/test_avg_pool2d_converter.py
@@ -1,4 +1,4 @@
-# Copyright 2024 NXP
+# Copyright 2024,2026 NXP
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
@@ -28,7 +28,10 @@
     ToNCHWPreprocess,
     ToNHWCPreprocess,
 )
+from executorch.backends.nxp.tests.graph_verifier import BaseGraphVerifier
 from executorch.backends.nxp.tests.models import AvgPool2dConvModule, AvgPool2dModule
+
+from executorch.backends.nxp.tests.nsys_testing import lower_run_compare
 from torch.export import ExportedProgram
 from executorch.backends.nxp.tests.use_qat import *  # noqa F403
 from executorch.exir.dialects._ops import ops as exir_ops
@@ -296,3 +299,73 @@ def test_from_avg_pool_1d(mocker):
         tflite_input_preprocess=ToChannelLastPreprocess(),
         tflite_output_preprocess=ToChannelFirstPreprocess(),
     )
+
+
+class TestAvgPool2DNewNeutronFlow:
+    def test__basic_nsys_inference(self):
+        input_shape = (2, 4, 6, 7)
+        model = AvgPool2dModule(False, 0)
+        graph_verifier = BaseGraphVerifier(
+            exp_num_delegate_call_nodes=1,  # Delegated AvgPool.
+            exp_non_delegated_nodes=[],
+        )
+
+        lower_run_compare(
+            model, input_shape, graph_verifier, use_new_flow_neutron_c=True
+        )
+
+    def test__kernel_size_limit(self):
+        kernel_size = (1, 4096)
+        input_shape = (1, 4) + kernel_size
+        model = AvgPool2dModule(False, 0, kernel_size)
+        graph_verifier = BaseGraphVerifier(
+            exp_num_delegate_call_nodes=1,  # Delegated AvgPool.
+            exp_non_delegated_nodes=[],
+        )
+
+        lower_run_compare(
+            model, input_shape, graph_verifier, use_new_flow_neutron_c=True
+        )
+
+    def test__kernel_size_limit_exceeded(self):
+        kernel_size = (1, 4097)  # Exceeds the kernel size limit.
+        input_shape = (1, 4) + kernel_size
+        model = AvgPool2dModule(False, 0, kernel_size)
+
+        delegated_ep = to_quantized_edge_program(
+            model, input_shape, use_new_flow_neutron_c=True
+        ).exported_program()
+
+        # Make sure the `avg_pool2d` was NOT delegated.
+        assert not graph_contains_any_of_ops(
+            delegated_ep.graph, [ExecutorchDelegateCall]
+        )
+        assert graph_contains_any_of_ops(delegated_ep.graph, [AvgPool2D])
+
+    def test__stride_limit(self):
+        stride = 4096
+        input_shape = (1, 4, 1, 4096)
+        model = AvgPool2dModule(False, 0, 1, stride)
+        graph_verifier = BaseGraphVerifier(
+            exp_num_delegate_call_nodes=1,  # Delegated AvgPool.
+            exp_non_delegated_nodes=[],
+        )
+
+        lower_run_compare(
+            model, input_shape, graph_verifier, use_new_flow_neutron_c=True
+        )
+
+    def test__stride_limit_exceeded(self):
+        stride = 4097  # Exceeds the stride limit.
+        input_shape = (1, 4, 1, 4096)
+        model = AvgPool2dModule(False, 0, 1, stride)
+
+        delegated_ep = to_quantized_edge_program(
+            model, input_shape, use_new_flow_neutron_c=True
+        ).exported_program()
+
+        # Make sure the `avg_pool2d` was NOT delegated.
+        assert not graph_contains_any_of_ops(
+            delegated_ep.graph, [ExecutorchDelegateCall]
+        )
+        assert graph_contains_any_of_ops(delegated_ep.graph, [AvgPool2D])
diff --git a/backends/nxp/tests/models.py b/backends/nxp/tests/models.py
index 17bea708352..c29491fa4e3 100644
--- a/backends/nxp/tests/models.py
+++ b/backends/nxp/tests/models.py
@@ -348,12 +348,12 @@ def forward(self, x):
 
 
 class AvgPool2dModule(torch.nn.Module):
-    def __init__(self, count_include_pad, padding=0):
+    def __init__(self, count_include_pad, padding=0, kernel_size=3, stride=2):
         super().__init__()
 
         self.avg_pool = torch.nn.AvgPool2d(
-            kernel_size=3,
-            stride=2,
+            kernel_size=kernel_size,
+            stride=stride,
             padding=padding,
             count_include_pad=count_include_pad,
         )