TensorRT icon indicating copy to clipboard operation
TensorRT copied to clipboard

feat: support aten.as_strided converter

Open chohk88 opened this issue 1 year ago • 0 comments

Description

New feature to support aten.as_strided converter converter. Our implementation focuses on accurately calculating the indices required

Fixes # (issue)

Example Usage

Given a tensor x:

x = torch.tensor([[1, 2, 3],
                  [4, 5, 6],
                  [7, 8, 9]])

Computing indices_tensor

Example 1: Stride=[1, 2]

  • Operation: torch.as_strided(x, (2, 2), (1, 2), 0)
  • indices_tensor: [0, 2, 1, 3] from stride [1, 2].
  • Output: [[1, 3], [2, 4]]

Example 2: Stride=[2, 1]

  • Operation: torch.as_strided(x, (2, 2), (2, 1), 0)
  • indices_tensor: [0, 1, 2, 3] from stride [2, 1].
  • Output: [[1, 2], [3, 4]]

as_strided converter calculates indices_tensor using size and stride, then reshapes the tensor accordingly

Type of change

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

Checklist:

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

chohk88 avatar Apr 09 '24 09:04 chohk88