Adding support for list of tensors as return type.
Nobody has claimed this yet.
- Dominant language
- C++
- Stars
- 1.9k
- Forks
- 736
- Avg merge
- 5d 22h
- Merged PRs (30d)
- 15
Description
Some of the operations like `aten.split`and `aten.split_with_sizes` return a list of tensors. This is currently not supported in torch-mlir.
```
def split_with_sizes(self: Tensor, split_sizes: List[int], dim: int = 0) -> List[Tensor]:
num_splits = len(split_sizes)
splits = []
start_idx = 0
for i in range(num_splits):
length = split_sizes[i]
splits.append(self.narrow(dim, start_idx, length))
start_idx += length
return splits
```
```
def split(self: Tensor, split_size: int, dim: int = 0) -> List[Tensor]:
input_sizes = self.shape
dim_size = input_sizes[dim]
if split_size == 0:
assert(dim_size == 0)
return [self]
chunks = (dim_size + split_size - 1) // split_size
split_sizes = [split_size for i in range(chunks)]
split_sizes[chunks - 1] = split_size - (split_size * chunks - dim_size)
return aten.split_with_sizes(self, split_sizes, dim)
```
I have started working on this. Please share some thoughts over this issue.
CC: @silvasean @cathyzhyi
Contributor guide
No contributing guide indexed for this repository
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start from the aten.split and aten.split_with_sizes definitions shown in the issue and trace how torch-mlir handles their List[Tensor] return types. No files or tests are named, so locate the return-type handling and relevant operation tests first. Done means both operations support list-of-tensor results without unsupported-type failures.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp, python
- Domain
- compilers
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 25/100