    def _infer_shapes(
        self,
        x: Tensor,
        start_indices: List[Union[IntVar, IntVarTensor, Optional[int]]],
        end_indices: List[Union[IntVar, IntVarTensor, Optional[int]]],
    ) -> List[IntVar]:
        """Infers shape for dynamic_slice."""
        # TODO: Handle start_indices/end_indices that are not int.

        x_shape = x._attrs["shape"]
        output_shape = []
        for dim_val, start, end in zip(x_shape, start_indices, end_indices):
            if start == 0 and end == MAX_INT32:
                # Slicing along the whole dim.
                output_shape.append(dim_val)
            elif isinstance(dim_val, IntImm):
                # Slicing a static dimension.
                start, end = dynamic_slice.normalize_start_end_indices(
                    dim_val.value(), start, end
                )
                output_shape.append(IntImm(end - start))
            elif start >= 0 and end >= 0:
                # Fixed size from start and end.
                output_shape.append(IntImm(end - start))
            else:
                output_shape.append(self._infer_dynamic_dim(dim_val, start, end))

        return output_shape
Infers shape for dynamic_slice.
python/aitemplate/compiler/ops/tensor/dynamic_slice.py