Skip to content

Commit

Permalink
removing chunk op, changing the skip test message and remove the dyna…
Browse files Browse the repository at this point in the history
…mic shape check of input in validator
  • Loading branch information
apbose committed Sep 4, 2024
1 parent 885046d commit b6851bd
Show file tree
Hide file tree
Showing 2 changed files with 5 additions and 29 deletions.
28 changes: 1 addition & 27 deletions py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py
Original file line number Diff line number Diff line change
Expand Up @@ -692,9 +692,7 @@ def aten_ops_softmax(

@dynamo_tensorrt_converter(
torch.ops.aten.split.Tensor,
capability_validator=(
has_static_shapes_in_args([0]) and has_static_shapes_in_args([1])
),
capability_validator=has_static_shapes_in_args([1]),
supports_dynamic_shapes=True,
)
@dynamo_tensorrt_converter(
Expand Down Expand Up @@ -905,30 +903,6 @@ def aten_ops_slice(
)


@dynamo_tensorrt_converter(torch.ops.aten.chunk.default)
@enforce_tensor_types(
{
0: (TRTTensor,),
}
)
def aten_ops_chunk(
ctx: ConversionContext,
target: Target,
args: Tuple[Argument, ...],
kwargs: Dict[str, Argument],
name: str,
) -> Union[TRTTensor, Sequence[TRTTensor]]:
return impl.slice.chunk(
ctx,
target,
SourceIR.ATEN,
name,
args[0],
args[1],
args_bounds_check(args, 2, 0),
)


@dynamo_tensorrt_converter(torch.ops.aten.cumsum.default, supports_dynamic_shapes=True)
@enforce_tensor_types(
{
Expand Down
6 changes: 4 additions & 2 deletions tests/py/dynamo/conversion/test_chunk_aten.py
Original file line number Diff line number Diff line change
Expand Up @@ -85,8 +85,10 @@ def forward(self, input):


#######################Dynamic cases#######################
#The tests are skipped for now. Will be addressed once https://github.com/pytorch/pytorch/issues/134663 is addressed
@unittest.skip("Pending aten.split converter. Currently tested by E2E")
# The tests are skipped for now. Will be addressed once https://github.com/pytorch/pytorch/issues/134663 is addressed
@unittest.skip(
"Pending aten.split dynamic input torch.export guard bug. Issue- pytorch/pytorch#134663"
)
class TestChunkDynamicConverter(DispatchTestCase):
@parameterized.expand(
[
Expand Down

0 comments on commit b6851bd

Please sign in to comment.