Skip to content

Support negative dimensions on "aten.split_with_sizes_copy.default" #16006

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
+3 −1
Open

Support negative dimensions on "aten.split_with_sizes_copy.default" #16006

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
5d58b93
Support negative dimensions on "aten.split_with_sizes_copy.default"
jgibson2 Nov 28, 2025
5f1a078
Report original uncorrected dimension on failure.
jgibson2 Nov 28, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion backends/apple/mps/operators/shape_ops.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -243,10 +243,12 @@ def define_node(
split_sizes = eval_shape(cast(torch.SymInt, node.args[1]))
dim = cast(int, node.args[2])
input_shape = get_shape(get_input_node(node, 0))
if dim < 0:
dim += len(input_shape)

if dim < 0 or dim >= len(input_shape):
Copy link

Copilot AI Nov 28, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The validation condition on line 249 will always fail for negative dimensions after normalization. If the original dim was less than -len(input_shape), the normalized value will still be negative, but this is the intended behavior to catch out-of-bounds dimensions. However, the current implementation normalizes first, then validates, which means valid negative dimensions are handled correctly but the error condition dim < 0 is now checking the normalized value. The validation should only check dim >= len(input_shape) after normalization, or validate the original dimension range before normalization.

Suggested change
if dim < 0 or dim >= len(input_shape):
if dim >= len(input_shape):

Copilot uses AI. Check for mistakes.
Copy link
Author

@jgibson2 jgibson2 Nov 28, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The negative dimension could be outside the correct range even after normalization (e.g. if there are two dimensions and the input dimension is -3). This should probably still be an error.

raise RuntimeError(
f"split_copy: dim {dim} out of range for input tensor with {len(input_shape)} dimensions"
f"split_copy: dim {cast(int, node.args[2])} out of range for input tensor with {len(input_shape)} dimensions"
)

mps_node = MPSNode(
Expand Down