Skip to content

Add support for SymInt start values in slice_scatter_decomposition #4185

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
shahe-nv wants to merge 1 commit into
base: main
Choose a base branch
from
+7 −2
Open

Add support for SymInt start values in slice_scatter_decomposition #4185

Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
9 changes: 7 additions & 2 deletions py/torch_tensorrt/dynamo/lowering/_decompositions.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -197,8 +197,13 @@ def slice_scatter_decomposition(
dim_size = input_tensor.shape[dim]
device_input_tensor = input_tensor.device

start = 0 if start is None else start # Ensure start is int
start = get_positive_dim(start, input_tensor.shape[dim])
if start is None:
start = 0
elif isinstance(start, int):
start = get_positive_dim(start, dim_size)
elif isinstance(start, torch.SymInt):
if start < 0:
start = start + dim_size
Comment on lines +200 to +206
Copy link
Copy Markdown
Collaborator

@zewenli98 zewenli98 Apr 30, 2026

Choose a reason for hiding this comment

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

According to the aten schema, all of start, end, and step could be SymInt, not just start. Can you fix all of them btw? I think a more general way is to support SymInt in get_positive_dim() like:

def get_positive_dim(
    dim: Union[int, SymInt, Sequence[int], Sequence[SymInt]], dim_size: int
):

if end is None: # Ensure end is int
Copy link
Copy Markdown
Collaborator

@apbose apbose Apr 14, 2026

Choose a reason for hiding this comment

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

Hi can you have a testcase for the above?

Copy link
Copy Markdown
Author

@shahe-nv shahe-nv Apr 14, 2026

Choose a reason for hiding this comment

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

Hi, I tested with my ZoomASR code. Without this, it will crash. After adding this it will run.

Copy link
Copy Markdown
Collaborator

@apbose apbose Apr 15, 2026

Choose a reason for hiding this comment

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

Thanks for the info! Generally for any change in the codebase, we recommend adding a corresponding test case to capture the corner case for future CI runs.

Copy link
Copy Markdown
Author

@shahe-nv shahe-nv Apr 27, 2026

Choose a reason for hiding this comment

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

Hi, do you know where I need to add the test? Do you have examples?

Copy link
Copy Markdown
Collaborator

@zewenli98 zewenli98 Apr 30, 2026

Choose a reason for hiding this comment

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

@shahe-nv I think @apbose meant you could add unit tests in tests/py/dynamo/lowering/test_decompositions.py

end = dim_size
end = (
Expand Down
Loading