wehub-resource-sync
144 lines
4.4 KiB
Python
144 lines
4.4 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
"""Indexing operators."""
|
|
|
|
from ..expr import Expr
|
|
from ..utils import convert_to_expr
|
|
from . import _ffi_api
|
|
|
|
PrimExprLike = int | Expr
|
|
|
|
|
|
def take(x: Expr, indices: Expr, axis: int | None = None, mode: str = "fast") -> Expr:
|
|
"""Take elements from a tensor along an axis.
|
|
Its semantic is mostly similar to `numpy.take`
|
|
(https://numpy.org/doc/stable/reference/generated/numpy.take.html),
|
|
which can cover `torch.take` (https://pytorch.org/docs/stable/generated/torch.take.html) and
|
|
`onnx.gather` (https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Gather-13).
|
|
|
|
Parameters
|
|
----------
|
|
x : relax.Expr
|
|
The source tensor.
|
|
|
|
indices : relax.Expr
|
|
The indices of the values to extract.
|
|
|
|
axis : Optional[int]
|
|
The axis over which to select values.
|
|
If it is none, the input tensor is required to be one-dimensional.
|
|
|
|
mode : str
|
|
Specifies how out-of-bounds indices will behave.
|
|
- fast (default): extra indices lead to seg fault (user must make sure indices are in-bound)
|
|
- nan: produce NaNs for out-of-bounds indices
|
|
- wrap: wrap around the indices
|
|
- clip: clip to the range
|
|
Returns
|
|
-------
|
|
ret : relax.Expr
|
|
The taken result.
|
|
"""
|
|
return _ffi_api.take(x, indices, axis, mode) # type: ignore
|
|
|
|
|
|
def strided_slice(
|
|
x: Expr,
|
|
axes: Expr,
|
|
begin: Expr,
|
|
end: Expr,
|
|
strides: Expr | None = None,
|
|
assume_inbound: bool = False,
|
|
) -> Expr:
|
|
"""Strided slice of a tensor.
|
|
|
|
Parameters
|
|
----------
|
|
x : relax.Expr
|
|
The source tensor to be sliced.
|
|
|
|
axes : List[int]
|
|
Axes along which slicing is applied.
|
|
|
|
begin : List[PrimExprLike]
|
|
The indices to begin with in the slicing, inclusive.
|
|
|
|
end : List[PrimExprLike]
|
|
The indices indicating end of the slice, exclusive.
|
|
|
|
strides : Optional[List[PrimExprLike]]
|
|
Specifies the stride values, it can be negative in that case,
|
|
the input tensor will be reversed in that particular axis.
|
|
If not specified, it by default is an list of ones of the same length as `axes`.
|
|
|
|
assume_inbound : bool
|
|
Whether to assume the indices are in bound. If it is set to false,
|
|
out of bound indices will be clipped to the bound.
|
|
Returns
|
|
-------
|
|
ret : relax.Expr
|
|
The sliced result.
|
|
|
|
Note
|
|
----
|
|
strided_slice require the input `begin`, `end` and `strides` to have the
|
|
same length as `axes`.
|
|
"""
|
|
axes = convert_to_expr(axes)
|
|
begin = convert_to_expr(begin)
|
|
end = convert_to_expr(end)
|
|
if strides is not None:
|
|
strides = convert_to_expr(strides)
|
|
return _ffi_api.strided_slice(x, axes, begin, end, strides, assume_inbound) # type: ignore
|
|
|
|
|
|
def dynamic_strided_slice(
|
|
x: Expr,
|
|
begin: Expr,
|
|
end: Expr,
|
|
strides: Expr,
|
|
) -> Expr:
|
|
"""Dynamic strided slice of a tensor. `begin`, `end`, `strides` can be computed at runtime.
|
|
|
|
Parameters
|
|
----------
|
|
x : Expr
|
|
The source tensor to be sliced.
|
|
|
|
begin : Expr
|
|
The indices to begin with in the slicing, inclusive.
|
|
|
|
end : Expr
|
|
The indices indicating end of the slice, exclusive.
|
|
|
|
strides : Expr
|
|
Specifies the stride values, it can be negative in that case,
|
|
the input tensor will be reversed in that particular axis.
|
|
If not specified, it by default is an list of ones of the same length as `axes`.
|
|
|
|
Returns
|
|
-------
|
|
ret : relax.Expr
|
|
The sliced result.
|
|
|
|
Note
|
|
----
|
|
dyn_strided_slice require the input `begin`, `end` and `strides` to have the
|
|
same length as rank of `data` tensor.
|
|
"""
|
|
return _ffi_api.dynamic_strided_slice(x, begin, end, strides) # type: ignore
|