Function strided_slice

python/tvm/topi/transform.py:176–235

Slice of an array. Parameters ---------- a : tvm.te.Tensor The tensor to be sliced. begin : list of int The indices to begin with in the slicing. end : list of int Indices indicating end of the slice. strides : list of int, optional Specifi

(a, begin, end, strides=None, axes=None, slice_mode="end", assume_inbound=True)

def strided_slice(a, begin, end, strides=None, axes=None, slice_mode="end", assume_inbound=True):
    """Slice of an array.

    Parameters
    ----------
    a : tvm.te.Tensor
        The tensor to be sliced.

    begin : list of int
        The indices to begin with in the slicing.

    end : list of int
        Indices indicating end of the slice.

    strides : list of int, optional
        Specifies the stride values, it can be negative
        in that case, the input tensor will be reversed
        in that particular axis.

    axes : list of int, optional
        Axes along which slicing is applied. When it is specified, begin, end
        strides, and axes need to a list of integers of the same length.

    slice_mode : str, optional
        The slice mode [end, size].
        end - The ending indices for the slice [default].
        size - The input strides will be ignored, input end in this mode indicates
        the sizeof a slice starting at the location specified by begin. If end[i]
        is -1, all remaining elements in that dimension are included in the slice.

    assume_inbound: bool, optional
        A flag to indicate if all indices are assumed to be inbound

    Returns
    -------
    ret : tvm.te.Tensor
    """
    if (
        isinstance(begin, tvm.te.Tensor)
        or isinstance(end, tvm.te.Tensor)
        or isinstance(strides, tvm.te.Tensor)
    ):
        assert axes is None, "axes argument is not supported by dynamic strided slice yet."
        if not isinstance(begin, tvm.te.Tensor):
            begin = const_vector(begin)
        if not isinstance(end, tvm.te.Tensor):
            end = const_vector(end)
        if strides is None:
            strides = [1] * begin.shape[0].value
        if not isinstance(strides, tvm.te.Tensor):
            strides = const_vector(strides)
        return cpp.dynamic_strided_slice(a, begin, end, strides)
    if strides is None:
        strides = []
    if axes is None:
        axes = []
    # axes is a list of host integers on the C++ side (Array<int64_t>); unwrap any
    # IntImm entries that callers may pass through (e.g. relax legalize pipeline).