NXP backend: added aten.split support #16490

novak-vaclav · 2026-01-07T16:05:50Z

Summary

adds support for aten.split operator

Test plan

tests can be manually run using pytest -c /dev/null backends/nxp/tests/

cc @robert-kalmar @JakeStevens @digantdesai @MartinPavella

pytorch-bot · 2026-01-07T16:05:55Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/16490

📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 2 New Failures, 1 Unrelated Failure

As of commit abb63f7 with merge base 0391fe7 ():

NEW FAILURES - The following jobs have failed:

pull / test-mediatek-models-linux / linux-job (gh)
RuntimeError: Command docker exec -t 5f3512b4026a4cd09b82eaac7b1cb8411797322a24d0e7b99deb7c1c11e5056d /exec failed with exit code 1
pull / unittest-nxp-neutron / linux-job (gh)
RuntimeError: Command docker exec -t 0b19b44f559e780226a23f36f24524288f5198e61ac979317ccefbfed78c1731 /exec failed with exit code 2

UNSTABLE - The following job is marked as unstable, possibly due to flakiness on trunk:

pull / android / run-emulator (gh) (#16137)
Timeout waiting for emulator to boot.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

novak-vaclav · 2026-01-07T16:06:08Z

@pytorchbot label "release notes: nxp"

novak-vaclav · 2026-01-07T16:06:12Z

@pytorchbot label "module: nxp"

novak-vaclav · 2026-01-07T16:08:18Z

Fixed all issues mentioned in this PR (#16276), including the comment of @roman-janik-nxp
But due to my error I had to create a new PR.

Copilot

Pull request overview

This PR adds support for the aten.split operator in the NXP backend by introducing a new pass that decomposes split operations into slice operations. The key addition is the DecomposeSplitToSlicesPass, which transforms aten.split.Tensor, aten.split.default, and aten.split_with_sizes.default operations into equivalent aten.slice.Tensor operations that the backend can handle.

Key changes:

Added DecomposeSplitToSlicesPass to decompose split operations into slices
Removed ConvertUnsqueezeToViewPass from the default pass manager
Added test models (SplitWithSize, SplitWithSections, GRUModel) to support split testing
Introduced comprehensive tests for the new decomposition pass

Reviewed changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 10 comments.

Show a summary per file

File	Description
backends/nxp/aten_passes/decompose_split_to_slices_pass.py	New pass that decomposes split operations into multiple slice operations for backend compatibility
backends/nxp/aten_passes/neutron_aten_pass_manager.py	Added DecomposeSplitToSlicesPass to default passes and removed ConvertUnsqueezeToViewPass import
backends/nxp/tests/test_decompose_split_to_slices.py	Comprehensive test suite for the new split decomposition pass
backends/nxp/tests/models.py	Added test models (GRUModel, SplitWithSize, SplitWithSections) to support split operation testing
backends/nxp/tests/test_split_group_convolution.py	Added workaround to exclude DecomposeSplitToSlicesPass from affecting group convolution tests

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

backends/nxp/tests/test_split_group_convolution.py

Copilot · 2026-01-07T16:13:46Z

backends/nxp/aten_passes/decompose_split_to_slices_pass.py

+        getitem_nodes = list(split_node.users.keys())
+        slice_nodes = []
+        for i in range(len(starts)):
+            slice_arguments = (input_node, dim, starts[i], ends[i])
+            with self.graph_module.graph.inserting_after(split_node):
+                slice_node = self._create_slice_node(*slice_arguments)
+                slice_nodes.append(slice_node)
+
+                getitem_node = getitem_nodes[i]
+                getitem_node.replace_all_uses_with(slice_node)
+
+                self.graph_module.graph.erase_node(getitem_node)


The code assumes that getitem_nodes from split_node.users.keys() maintains the correct order corresponding to split indices (0, 1, 2, ...). While dictionaries in Python 3.7+ maintain insertion order, the users of a split node may not be ordered by their getitem indices. This could lead to incorrect replacement where slice_node for index i is paired with the wrong getitem_node. Consider explicitly matching getitem nodes to their indices, or sorting them by their getitem index argument to ensure correct correspondence.

Copilot · 2026-01-07T16:13:47Z

backends/nxp/aten_passes/neutron_aten_pass_manager.py

+from executorch.backends.nxp.aten_passes.decompose_split_to_slices_pass import (
+    DecomposeSplitToSlicesPass,
 )


Removing the import of ConvertUnsqueezeToViewPass from this module may break existing code that imports it from neutron_aten_pass_manager. Consider either:

Keeping the import and re-exporting it for backward compatibility, even if it's not used in the default pass list

Adding an explicit __all__ list that includes both ConvertUnsqueezeToViewPass and DecomposeSplitToSlicesPass for users who want to create custom pass managers

Updating all references to import from convert_unsqueeze_to_view directly

Note that test_convert_unsqueeze_to_view.py currently imports this class from this module and will fail after this change.

Copilot · 2026-01-07T16:13:47Z

backends/nxp/tests/test_decompose_split_to_slices.py

+        pytest.param((4, 8), [3, 3, 2], 1, id="2D."),
+    ],
+)
+def test_decompose_split_with_section(mocker, input_shape, sections, dim):


The mocker parameter is declared but never used in this test function. Consider removing it to keep the test signature clean and avoid confusion.

Copilot · 2026-01-07T16:13:47Z

backends/nxp/tests/test_decompose_split_to_slices.py

+        pytest.param(2, id="2 GRU layers"),
+    ],
+)
+def test_decompose_gru_with_split(mocker, gru_layers):


The mocker parameter is declared but never used in this test function. Consider removing it to keep the test signature clean and avoid confusion.

Suggested change

def test_decompose_gru_with_split(mocker, gru_layers):

def test_decompose_gru_with_split(gru_layers):

Copilot · 2026-01-07T16:13:48Z

backends/nxp/aten_passes/neutron_aten_pass_manager.py

        passes: list[PassType] = passes or [
+            DecomposeSplitToSlicesPass(),
            FuseBatchNormWithConvPass(),
            FuseBatchNormWithLinearPass(),
            SplitGroupConvolution(),
            SplitGRUBasedOnNumLayers(),
            RemoveNodesWithKnownOutputs(),
            FuseLinearAndAddPass(),
            MoveActivationBeforeConcat(neutron_target_spec),
-            ConvertUnsqueezeToViewPass(),
        ]


The pass ordering may be problematic. DecomposeSplitToSlicesPass is placed first in the pass list, but SplitGroupConvolution (at position 4) creates new aten.split.default operations that won't be decomposed. This creates an inconsistency where:

Pre-existing splits in the model will be decomposed to slices

Splits created by SplitGroupConvolution will remain as splits

Consider either:

Moving DecomposeSplitToSlicesPass to run after SplitGroupConvolution if all splits should be decomposed

Keeping the current order if splits from group convolution decomposition should intentionally remain as splits for backend delegation

Making the pass more selective to skip splits that are part of group convolution patterns

The test workaround in test_split_group_convolution.py suggests this ordering issue is known but unresolved.

Copilot · 2026-01-07T16:13:48Z

backends/nxp/aten_passes/decompose_split_to_slices_pass.py

+            split_nodes_chunks = list(split_nodes_chunks)
+
+        if not isinstance(split_nodes_chunks, list):
+            raise RuntimeError("Faulty split chunks")


The error message "Faulty split chunks" is not descriptive enough for debugging. Consider providing more context about what went wrong, such as the actual type received and what was expected, to help developers diagnose issues more quickly.

Suggested change

raise RuntimeError("Faulty split chunks")

raise RuntimeError(

f"Faulty split chunks: expected a list of chunk tensors (or a tuple "

f"convertible to list) in meta['val'], but got "

f"{type(split_nodes_chunks).__name__} with value "

f"{split_nodes_chunks!r} for node {split_node!r}"

)

Copilot · 2026-01-07T16:13:48Z

backends/nxp/tests/test_decompose_split_to_slices.py

+        pytest.param((4, 8), 5, 1, id="2D."),
+    ],
+)
+def test_decompose_split_with_size(mocker, input_shape, split_size, dim):


The mocker parameter is declared but never used in this test function. Consider removing it to keep the test signature clean and avoid confusion.

Suggested change

def test_decompose_split_with_size(mocker, input_shape, split_size, dim):

def test_decompose_split_with_size(input_shape, split_size, dim):

Copilot · 2026-01-07T16:13:49Z

backends/nxp/tests/test_decompose_split_to_slices.py

+        pytest.param((8, 4), [4], 1, id="2D, one chunk using sections."),
+    ],
+)
+def test_decompose_split_with_one_chunk(mocker, input_shape, size_or_sections, dim):


The mocker parameter is declared but never used in this test function. Consider removing it to keep the test signature clean and avoid confusion.

Suggested change

def test_decompose_split_with_one_chunk(mocker, input_shape, size_or_sections, dim):

def test_decompose_split_with_one_chunk(input_shape, size_or_sections, dim):

Copilot · 2026-01-07T16:13:49Z

backends/nxp/aten_passes/decompose_split_to_slices_pass.py

+                getitem_node = list(split_node.users)[0]
+                getitem_node.replace_all_uses_with(input_node)
+
+                self.graph_module.graph.erase_node(getitem_node)
+                self.graph_module.graph.erase_node(split_node)
+
+                made_changes = True
+                continue


When handling single-chunk splits, the code assumes there is exactly one user (getitem node) by using list(split_node.users)[0]. This could raise an IndexError if the split has no users, or silently ignore additional users if there are multiple. Consider adding validation to ensure there is exactly one user, or handling edge cases where there are zero or multiple users.

Suggested change

getitem_node = list(split_node.users)[0]

getitem_node.replace_all_uses_with(input_node)

self.graph_module.graph.erase_node(getitem_node)

self.graph_module.graph.erase_node(split_node)

made_changes = True

continue

split_users = list(split_node.users)

# Only perform this optimization when there is exactly one user.

# If there are zero or multiple users, fall back to the generic

# slice-based decomposition below to avoid unsafe rewrites.

if len(split_users) == 1:

getitem_node = split_users[0]

getitem_node.replace_all_uses_with(input_node)

self.graph_module.graph.erase_node(getitem_node)

self.graph_module.graph.erase_node(split_node)

made_changes = True

continue

MartinPavella · 2026-01-08T16:16:46Z

backends/nxp/tests/test_split_group_convolution.py

@@ -1,4 +1,4 @@
-# Copyright 2025 NXP
+# Copyright 2026 NXP


When updating an old file, please keep the original copyright year as well. In this case, please use:
# Copyright 2025-2026 NXP

MartinPavella · 2026-01-08T16:19:08Z

backends/nxp/tests/test_split_group_convolution.py


-        modified_module = NeutronAtenPassManager(neutron_target_spec)(
+        neutron_pass_manager = NeutronAtenPassManager(neutron_target_spec)
+        # Provisional solution to manually remove DecomposeSplitToSlicesPass until NeutronAtenPassManager is fixed


This is not true.
Regardless of how the NeutronAtenPassManager works, the pass is removed here not because of any bugs, but only because it would (correctly) modify the graph, in a way that would require changes to this test, and it would reduce clarity.

So the pass is removed only for testing purposes.

Even when we update the NeutronAtenPassManager as we discussed, this test will still have the pass removed.

backends/nxp/tests/test_split_group_convolution.py

MartinPavella · 2026-01-08T16:25:37Z

backends/nxp/aten_passes/neutron_aten_pass_manager.py

@@ -1 +1 @@
 # Copyright 2026 NXP


feat: added support for split operator

abb63f7

novak-vaclav requested a review from robert-kalmar as a code owner January 7, 2026 16:05

Copilot AI review requested due to automatic review settings January 7, 2026 16:05

meta-cla bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jan 7, 2026

pytorch-bot bot added release notes: nxp Changes to the NXP Neutron backend delegate module: nxp Issues related to NXP Neutron NPU delegation and code under backends/nxp/ labels Jan 7, 2026

Copilot started reviewing on behalf of novak-vaclav January 7, 2026 16:06 View session

Copilot AI reviewed Jan 7, 2026

View reviewed changes

MartinPavella requested changes Jan 8, 2026

View reviewed changes

	def test_decompose_gru_with_split(mocker, gru_layers):
	def test_decompose_gru_with_split(gru_layers):

-            raise RuntimeError("Faulty split chunks")
+            raise RuntimeError(
+                f"Faulty split chunks: expected a list of chunk tensors (or a tuple "
+                f"convertible to list) in meta['val'], but got "
+                f"{type(split_nodes_chunks).__name__} with value "
+                f"{split_nodes_chunks!r} for node {split_node!r}"
+            )

	def test_decompose_split_with_size(mocker, input_shape, split_size, dim):
	def test_decompose_split_with_size(input_shape, split_size, dim):

	def test_decompose_split_with_one_chunk(mocker, input_shape, size_or_sections, dim):
	def test_decompose_split_with_one_chunk(input_shape, size_or_sections, dim):

NXP backend: added aten.split support #16490

Are you sure you want to change the base?

NXP backend: added aten.split support #16490

Conversation

novak-vaclav commented Jan 7, 2026 • edited by pytorch-bot bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Test plan

Uh oh!

pytorch-bot bot commented Jan 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/16490

❌ 2 New Failures, 1 Unrelated Failure

Uh oh!

novak-vaclav commented Jan 7, 2026

Uh oh!

novak-vaclav commented Jan 7, 2026

Uh oh!

novak-vaclav commented Jan 7, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Copilot AI Jan 7, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jan 7, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jan 7, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jan 7, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jan 7, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jan 7, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jan 7, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jan 7, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jan 7, 2026

Choose a reason for hiding this comment

Uh oh!

MartinPavella Jan 8, 2026

Choose a reason for hiding this comment

Uh oh!

MartinPavella Jan 8, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

MartinPavella Jan 8, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

novak-vaclav commented Jan 7, 2026 •

edited by pytorch-bot bot

Loading

pytorch-bot bot commented Jan 7, 2026 •

edited

Loading