TensorRT
TensorRT copied to clipboard
feat: support aten.as_strided converter
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