removing chunk op, changing the skip test message and remove the dynamic shape check of input in validator

Author: apbose

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -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(
@@ -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(
     {

py/torch_tensorrt/dynamo/conversion/impl/slice/ops.py

@@ -324,61 +324,6 @@ def expand(
     return layer.get_output(0)
 
 
-def chunk(
-    ctx: ConversionContext,
-    target: Target,
-    source_ir: Optional[SourceIR],
-    name: str,
-    input: TRTTensor,
-    chunks: int,
-    dim: int,
-) -> TRTTensor:
-    if chunks <= 0:
-        raise RuntimeError(
-            f"chunk expects `chunks` to be greater than 0, got: {chunks}"
-        )
-
-    shape = input.shape
-    dim = get_positive_dim(dim, len(shape))
-
-    if dim >= len(shape):
-        raise RuntimeError(
-            f"chunk expects `dim` to be less than the length of input shape, got: {dim}"
-        )
-
-    dynamic_shape = has_dynamic_shape(input.shape)
-    if dynamic_shape > 0:
-        # Check whether slice target dim is dynamic shape dim
-        assert input.shape[dim] != -1, "Can't chunk on dynamic shape dimension!"
-
-    size_dim = shape[dim]
-    chunk_size = math.ceil(size_dim / chunks)
-    result = []
-    start = 0
-    end = min(start + chunk_size, size_dim)
-    cnt = 0
-
-    while start < end:
-        result.append(
-            slice_op(
-                ctx,
-                target,
-                source_ir,
-                f"{name}_slice_{cnt}",
-                input,
-                dim,
-                start,
-                end,
-                1,
-            )
-        )
-        start = end
-        end = min(start + chunk_size, size_dim)
-        cnt += 1
-
-    return result
-
-
 def cumsum(
     ctx: ConversionContext,
     target: Target,

tests/py/dynamo/conversion/test_chunk_aten.py

@@ -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- https://github.com/pytorch/pytorch/issues/134663"
+)
 class TestChunkDynamicConverter(DispatchTestCase):
     @parameterized.expand(
         [