Skip to content

[ET-VK] Add aten.index.Tensor op and tests#17710

Merged
SS-JIA merged 1 commit intogh/SS-JIA/452/basefrom
gh/SS-JIA/452/head
Feb 25, 2026
Merged

[ET-VK] Add aten.index.Tensor op and tests#17710
SS-JIA merged 1 commit intogh/SS-JIA/452/basefrom
gh/SS-JIA/452/head

Conversation

@SS-JIA
Copy link
Copy Markdown
Contributor

@SS-JIA SS-JIA commented Feb 25, 2026

Stack from ghstack (oldest at bottom):

Implements aten.index.Tensor for the Vulkan backend, supporting 1D
self tensors with exactly one non-None index tensor. Includes buffer
and texture GLSL shaders, C++ operator registration, and correctness
tests.

Also extends the op test code generators to handle
c10::List<::std::optional<at::Tensor>> (Tensor?[] in ATen), which
is the C++ type for the indices argument of aten.index.Tensor.

Differential Revision: D94364638

Implements `aten.index.Tensor` for the Vulkan backend, supporting 1D
self tensors with exactly one non-None index tensor. Includes buffer
and texture GLSL shaders, C++ operator registration, and correctness
tests.

Also extends the op test code generators to handle
`c10::List<::std::optional<at::Tensor>>` (`Tensor?[]` in ATen), which
is the C++ type for the `indices` argument of `aten.index.Tensor`.

Differential Revision: [D94364638](https://our.internmc.facebook.com/intern/diff/D94364638/)

[ghstack-poisoned]
@pytorch-bot
Copy link
Copy Markdown

pytorch-bot Bot commented Feb 25, 2026

🔗 Helpful Links

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

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

❌ 3 New Failures

As of commit f306d9a with merge base 63f9724 (image):

NEW FAILURES - The following jobs have failed:

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

SS-JIA pushed a commit that referenced this pull request Feb 25, 2026
Implements `aten.index.Tensor` for the Vulkan backend, supporting 1D
self tensors with exactly one non-None index tensor. Includes buffer
and texture GLSL shaders, C++ operator registration, and correctness
tests.

Also extends the op test code generators to handle
`c10::List<::std::optional<at::Tensor>>` (`Tensor?[]` in ATen), which
is the C++ type for the `indices` argument of `aten.index.Tensor`.

Differential Revision: [D94364638](https://our.internmc.facebook.com/intern/diff/D94364638/)

ghstack-source-id: 344667758
Pull Request resolved: #17710
@meta-cla 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 Feb 25, 2026
@github-actions
Copy link
Copy Markdown

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

ValueListPtr indices_list = graph.get_value_list(indices_list_ref);
VK_CHECK_COND(
indices_list->size() == 1,
"index.Tensor: only one index tensor is supported");
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

How about Bool vs Int? I think you are implementing the integer case only, right? If that's the case, also add a VK_CHECK_COND saying that only integer index is supported.

@SS-JIA SS-JIA merged commit 5638f4b into gh/SS-JIA/452/base Feb 25, 2026
202 of 211 checks passed
@SS-JIA SS-JIA deleted the gh/SS-JIA/452/head branch February 25, 2026 19:12
@SS-JIA SS-JIA temporarily deployed to cherry-pick-bot February 25, 2026 19:12 — with GitHub Actions Inactive
SS-JIA pushed a commit that referenced this pull request Feb 25, 2026
Implements `aten.index.Tensor` for the Vulkan backend, supporting 1D
self tensors with exactly one non-None index tensor. Includes buffer
and texture GLSL shaders, C++ operator registration, and correctness
tests.

Also extends the op test code generators to handle
`c10::List<::std::optional<at::Tensor>>` (`Tensor?[]` in ATen), which
is the C++ type for the `indices` argument of `aten.index.Tensor`.

Differential Revision: [D94364638](https://our.internmc.facebook.com/intern/diff/D94364638/)

ghstack-source-id: 344667758
Pull Request resolved: #17710
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants