Skip to content

Conversation

ruoqianguo
Copy link
Contributor

Signed-off-by: Ruoqian Guo [email protected]

Description

Support multiple indices for aten::index.Tensor.
In Pytorch, aten::index.Tensor can receive several indices, like this:

x = torch.rand(5, 7, 10)
idx0 = torch.tensor([0, 0, 0, 0])
idx1 = torch.tensor([1, 1, 1, 1])
idx2 = torch.tensor([2, 2, 2, 2])
y1 = x[idx0, idx1, idx2] # result is like tensor([x[0, 1, 2], x[0, 1, 2], x[0, 1, 2], x[0, 1, 2]])
y2 = x[idx0] # result is like tensor([x[0, :, :], x[0, :, :], x[0, :, :], x[0, :, :]])

Previous, we only support one indice in PR #921 . Now we support serval indices for aten::index.Tensor.

Fixes #1274

Type of change

Please delete options that are not relevant and/or add your own.

  • New feature (non-breaking change which adds functionality)

Checklist:

  • My code follows the style guidelines of this project (You can use the linters)
  • I have performed a self-review of my own code
  • I have commented my code, particularly in hard-to-understand areas and hacks
  • I have made corresponding changes to the documentation
  • I have added tests to verify my fix or my feature
  • New and existing unit tests pass locally with my changes
  • I have added the relevant labels to my PR in so that relevant reviewers are notified

@github-actions github-actions bot added component: conversion Issues re: Conversion stage component: converters Issues re: Specific op converters component: core Issues re: The core compiler component: evaluators Issues re: Specific op evaluators component: tests Issues re: Tests labels Aug 25, 2022
Copy link
Collaborator

@narendasan narendasan left a comment

Choose a reason for hiding this comment

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

LGTM

@narendasan narendasan merged commit 22c0e17 into pytorch:master Aug 25, 2022
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
cla signed component: conversion Issues re: Conversion stage component: converters Issues re: Specific op converters component: core Issues re: The core compiler component: evaluators Issues re: Specific op evaluators component: tests Issues re: Tests
Projects
None yet
3 participants