Adding add_slice function in operator.py

Author: apbose

py/torch_tensorrt/fx/converters/operator.py

@@ -1128,3 +1128,51 @@ def add_expand(network, target, kwargs, name):
     layer = network.add_slice(input_val, start=start, shape=shape, stride=stride)
     set_layer_name(layer, target, name)
     return layer.get_output(0)
+
+
+def add_slice(network, target, kwargs, name):
+    input_val = kwargs["input"]
+    if not isinstance(input_val, TRTTensor):
+        raise RuntimeError(
+            f"slice_tensor received input {input_val} that is not part "
+            "of the TensorRT region!"
+        )
+
+    ranks = len(input_val.shape) + (1 if network.has_implicit_batch_dimension else 0)
+    dim = get_positive_dim(cast(int, kwargs["dim"]), ranks)
+    dynamic_shape = has_dynamic_shape(input_val.shape)
+    if network.has_implicit_batch_dimension:
+        if dim == 0:
+            raise RuntimeError(
+                f"We do not support slice_tensor at batch dim when it's implicit, got {dim}!"
+            )
+        dim = dim - 1
+    else:
+        if dynamic_shape:
+            # Check whether slice target dim is dynamic shape dim
+            assert input_val.shape[dim] != -1, "Can't chunk on dynamic shape dimension!"
+
+    start_int = cast(int, kwargs["start"])
+    stop_int = cast(int, kwargs["stop"])
+    step_int = cast(int, kwargs["step"])
+    start = [0] * len(input_val.shape)
+    start[dim] = start_int
+    stride = [1] * len(start)
+    stride[dim] = step_int
+    output_shape = list(input_val.shape)
+    output_shape[dim] = (stop_int - start_int) // step_int
+
+    if dynamic_shape > 0:
+        output_shape = get_shape_with_dynamic_shape(
+            network, output_shape, input_val, target, name
+        )
+    layer = network.add_slice(
+        input_val,
+        start=start,
+        shape=[] if dynamic_shape else output_shape,
+        stride=stride,
+    )
+    if dynamic_shape:
+        layer.set_input(2, output_shape)
+    set_layer_name(layer, target, name)
+    return layer.get_output(0)