aten::split converter #2232
Merged
aten::split converter #2232
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
ee21e73 to
0c53e8c
Compare
0c53e8c to
7be7512
Compare
7be7512 to
a332a44
Compare
gs-olive
reviewed
Aug 15, 2023
| @@ -279,6 +279,22 @@ def aten_ops_unsqueeze( | |||
| ) | |||
|
|
|||
|
|
|||
| @dynamo_tensorrt_converter( | |||
| torch.ops.aten.split.default, capability_validator=dynamic_unsupported | |||
There was a problem hiding this comment.
Switch to `torch.ops.aten.split.Tensor
| kwargs: Dict[str, Argument], | ||
| name: str, | ||
| ) -> Union[TRTTensor, Sequence[TRTTensor]]: | ||
| return impl.split(network, target, SourceIR.ATEN, name, args[0], args[1], args[2]) |
There was a problem hiding this comment.
Switch to impl.split.split
| source_ir: Optional[SourceIR], | ||
| name: str, | ||
| input: TRTTensor, | ||
| split_size_or_sections: Union[int, List(int)], |
There was a problem hiding this comment.
Replace List(int) with List[int]
| offset = 0 | ||
|
|
||
| if len(split_sizes) == 1: | ||
| num_splits = input.shape[dim] + split_sizes[0] - 1 // split_sizes[0] |
There was a problem hiding this comment.
Parenthesize first argument (num_splits = input.shape[dim] + split_sizes[0] - 1) // split_sizes[0]
| name: str, | ||
| input: TRTTensor, | ||
| split_size_or_sections: Union[int, List(int)], | ||
| dim: Optional[Any] = 0, |
There was a problem hiding this comment.
Replace with dim: int = 0, since this is a mandatory argument, as per native_functions.yaml.
Comment on lines
33
to
35
| if network.has_implicit_batch_dimension: | ||
| assert dim != 0, "Can't split on batch dim when it's implicit!" | ||
| dim -= 1 |
There was a problem hiding this comment.
Can remove this conditional, since implicit_batch_dimension is deprecated
| assert input.shape[dim] != -1, "Can't chunk on dynamic shape dimension!" | ||
|
|
||
| split_sizes = [] | ||
| if type(split_size_or_sections) == int: |
There was a problem hiding this comment.
Replace with isinstance(split_size_or_sections, int)
a332a44 to
3c09956
Compare
gs-olive
reviewed
Sep 19, 2023
| @@ -351,6 +351,34 @@ def aten_ops_softmax( | |||
| ) | |||
|
|
|||
|
|
|||
| @dynamo_tensorrt_converter( | |||
| torch.ops.aten.split.Tensor, capability_validator=dynamic_unsupported_with_args([1]) | |||
There was a problem hiding this comment.
The dynamic_unsupported_with_args import may also have been removed in the rebase
8b1d0d5 to
d74c61f
Compare
d74c61f to
4fed932
Compare
- Add support for selecting individual argument positions to check and expand checking to include symbolic types, which are sometimes passed in as arguments
gs-olive
approved these changes
Sep 21, 2023
| @@ -0,0 +1,81 @@ | |||
| from typing import Any, Dict, List, Optional, Sequence, Tuple, Union, cast | |||
There was a problem hiding this comment.
Minor point, but the linting may complain about unused import for cast here
5d1309a to
797d510
Compare
c3341c6 to
e95e0e4
Compare
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
aten::split converter