wehub-resource-sync
907 lines
26 KiB
Python
907 lines
26 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.
|
|
"""Manipulation operators."""
|
|
|
|
from collections.abc import Callable
|
|
|
|
from tvm.ir import is_prim_expr
|
|
from tvm.runtime import DataTypeCode
|
|
from tvm.tirx import FloatImm, IndexMap, IntImm
|
|
|
|
from ..expr import Expr, ShapeExpr, prim_value
|
|
from ..expr import Tuple as RxTuple
|
|
from . import _ffi_api
|
|
|
|
PrimExprLike = int | Expr
|
|
|
|
|
|
def broadcast_to(x: Expr, shape: tuple[PrimExprLike] | Expr) -> Expr:
|
|
"""Broadcasts a tensor to a specified shape.
|
|
|
|
Parameters
|
|
----------
|
|
x : relax.Expr
|
|
The input data to the operator.
|
|
|
|
shape : Union[Tuple[PrimExprLike], Expr]
|
|
The target shape.
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The broadcasted tensor.
|
|
"""
|
|
if isinstance(shape, tuple | list):
|
|
shape = ShapeExpr(shape)
|
|
return _ffi_api.broadcast_to(x, shape) # type: ignore
|
|
|
|
|
|
def concat(tensors: Expr | list[Expr], axis: int | None = 0) -> Expr:
|
|
"""Concatenate the input tensors along the given axis.
|
|
|
|
Parameters
|
|
----------
|
|
tensors : Union[relax.Expr, List[relax.Expr]]
|
|
An Expr in Tuple type, containing the tensors to be concatenated,
|
|
or a list of Tensors.
|
|
|
|
axis : Optional[int]
|
|
The axis along which the tensors are concatenated.
|
|
If `axis` is `None`, the input tensor is required to be flattened before concatenation.
|
|
|
|
Returns
|
|
-------
|
|
result: relax.Expr
|
|
The concatenated tensor.
|
|
"""
|
|
if isinstance(tensors, list | tuple):
|
|
tensors = RxTuple(tensors)
|
|
return _ffi_api.concat(tensors, axis) # type: ignore
|
|
|
|
|
|
def expand_dims(x: Expr, axis: int | list[int]) -> Expr:
|
|
"""Insert new axes at the positions given by `axis`.
|
|
|
|
Parameters
|
|
----------
|
|
x : relax.Expr
|
|
The input data to the operator.
|
|
|
|
axis : Union[int, List[int]]
|
|
The axes at which the input array are expanded.
|
|
All values are required to lie in range `[-data.ndim - 1, data.ndim]`, with the convention
|
|
of negative indexing.
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The transformed result.
|
|
"""
|
|
if isinstance(axis, int):
|
|
axis = [axis]
|
|
return _ffi_api.expand_dims(x, axis) # type: ignore
|
|
|
|
|
|
def flatten(x: Expr) -> Expr:
|
|
"""Flatten all the tensor dimensions into one.
|
|
|
|
Parameters
|
|
----------
|
|
x : relax.Expr
|
|
The input data to the operator.
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The flattened result.
|
|
"""
|
|
return _ffi_api.flatten(x) # type: ignore
|
|
|
|
|
|
def layout_transform(
|
|
x: Expr,
|
|
index_map: Callable | IndexMap,
|
|
pad_value: int | float | Expr | None = None,
|
|
axis_separators: int | str | None = None, # str for IndexMap.AXIS_SEPARATOR
|
|
input_axis_separators: int | str | None = None, # str for IndexMap.AXIS_SEPARATOR
|
|
):
|
|
"""Modifies the layout of a tensor.
|
|
|
|
Parameters
|
|
----------
|
|
x : relax.Expr
|
|
The input tensor to the operator.
|
|
|
|
index_map : Callable | IndexMap
|
|
The transformation to apply.
|
|
|
|
pad_value : Optional[int | float | Expr]
|
|
The value used for padding if the transformation results in implicit padding.
|
|
If not specified, any value can be used.
|
|
|
|
axis_separators : Optional[int | IndexMap.AXIS_SEPARATOR]
|
|
The axis_separators for index_map to create non flat buffers.
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The transformed tensor.
|
|
"""
|
|
default_index_dtype = "int64"
|
|
|
|
if callable(index_map):
|
|
index_map = IndexMap.from_func(index_map, index_dtype=default_index_dtype)
|
|
x_dtype = x.ty.dtype
|
|
|
|
# Explicitly convert python int/float pad_value to the x's type. If the default behavior
|
|
# is applied, it would be converted to int32/float32, which may not match the x's type.
|
|
if pad_value is None:
|
|
pass
|
|
elif not is_prim_expr(pad_value):
|
|
if x_dtype.matches_code(DataTypeCode.INT, DataTypeCode.UINT) and isinstance(pad_value, int):
|
|
pad_value = IntImm(x_dtype.dtype, pad_value)
|
|
elif x_dtype.matches_code(DataTypeCode.FLOAT, DataTypeCode.BFLOAT) and (
|
|
isinstance(pad_value, int | float)
|
|
):
|
|
pad_value = FloatImm(x_dtype.dtype, float(pad_value))
|
|
pad_value = prim_value(pad_value)
|
|
|
|
if axis_separators is None:
|
|
axis_separators = []
|
|
|
|
if input_axis_separators is None:
|
|
input_axis_separators = []
|
|
|
|
return _ffi_api.layout_transform(
|
|
x, index_map, pad_value, axis_separators, input_axis_separators
|
|
)
|
|
|
|
|
|
def permute_dims(x: Expr, axes: list[int] | None = None) -> Expr:
|
|
"""Permutes the dimensions of an array.
|
|
|
|
Parameters
|
|
----------
|
|
x : relax.Expr
|
|
The input data to the operator.
|
|
|
|
axes : Optional[List[int]]
|
|
The target axes order. If not specified, permute_dims will reverse the order of all axes.
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The transposed result.
|
|
"""
|
|
return _ffi_api.permute_dims(x, axes) # type: ignore
|
|
|
|
|
|
def reshape(x: Expr, shape: tuple[PrimExprLike] | Expr) -> Expr:
|
|
"""Reshape the input array.
|
|
|
|
``-1`` infers the dimension of the output shape by using the remainder of
|
|
the input dimensions keeping the size of the new array same as that of the input array.
|
|
At most one dimension of shape can be -1.
|
|
|
|
.. code-block:: python
|
|
|
|
x.shape = (2, 3, 4), shape = (6, 1, -1), result.shape = (6, 1, 4)
|
|
x.shape = (2, 3, 4), shape = (3, -1, 8), result.shape = (3, 1, 8)
|
|
x.shape = (2, 3, 4), shape = (-1,), result.shape = (24,)
|
|
|
|
Parameters
|
|
----------
|
|
x : relax.Expr
|
|
The input data to the operator.
|
|
|
|
shape : Union[Tuple[PrimExprLike], Expr]
|
|
The new shape. Should be compatible with the original shape.
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The reshaped result.
|
|
|
|
Note
|
|
----
|
|
The ``-1`` inference is only performed at compile-time.
|
|
That is to say, in any case the dimension length of ``-1`` cannot be inferred in
|
|
compile-time, an error will be thrown.
|
|
"""
|
|
if not isinstance(shape, tuple | list | Expr) or is_prim_expr(shape):
|
|
raise TypeError("shape must be a tuple/list or a Relax shape expression")
|
|
return _ffi_api.reshape(x, shape) # type: ignore
|
|
|
|
|
|
def split(
|
|
x: Expr,
|
|
indices_or_sections: int | list[PrimExprLike],
|
|
axis: int = 0,
|
|
) -> Expr:
|
|
"""Split input tensor along axis by sections or indices.
|
|
|
|
If indices_or_sections is an integer, the input will be divided equally
|
|
along given axis (if possible). Last section will be smaller if the tensor
|
|
size along the given dimension is not divisible by the integer.
|
|
|
|
If indices_or_sections is a tuple of mixture of int or Expr,
|
|
the entries indicate the indices where along axis the array is split.
|
|
|
|
Parameters
|
|
----------
|
|
x : relax.Expr
|
|
The tensor to be split.
|
|
|
|
indices_or_sections : Union[int, List[PrimExprLike]]
|
|
Indices or sections to split into. Accepts an int or a list.
|
|
|
|
axis : int
|
|
The axis over which to split.
|
|
|
|
Returns
|
|
-------
|
|
ret : relax.Expr
|
|
The computed result.
|
|
"""
|
|
if isinstance(indices_or_sections, int):
|
|
indices_or_sections = IntImm("int64", indices_or_sections)
|
|
return _ffi_api.split(x, indices_or_sections, axis) # type: ignore
|
|
|
|
|
|
def squeeze(x: Expr, axis: int | list[int] | None = None) -> Expr:
|
|
"""Squeeze axes in the array.
|
|
|
|
Parameters
|
|
----------
|
|
x : relax.Expr
|
|
The input data to the operator.
|
|
|
|
axis : Optional[Union[int, List[int]]
|
|
The set of axes to remove.
|
|
If axis = None, remove all axis of dimensions 1.
|
|
If any specified axis has dimension that does not equal 1, it is an error.
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The squeezed result.
|
|
"""
|
|
if isinstance(axis, int):
|
|
axis = [axis]
|
|
return _ffi_api.squeeze(x, axis) # type: ignore
|
|
|
|
|
|
def stack(tensors: Expr | list[Expr], axis: int = 0) -> Expr:
|
|
"""Stack the input tensors along a new axis.
|
|
|
|
Parameters
|
|
----------
|
|
tensors : Union[relax.Expr, List[relax.Expr]]
|
|
An Expr in Tuple type, containing the tensors to be stacked,
|
|
or a list of Tensors. All input tensors must have the same shape.
|
|
|
|
axis : int
|
|
The axis in the resulting tensor along which the input tensors will be stacked.
|
|
Negative values wrap around. Default is 0.
|
|
|
|
Returns
|
|
-------
|
|
result: relax.Expr
|
|
The stacked tensor with an additional dimension compared to the input tensors.
|
|
|
|
"""
|
|
if isinstance(tensors, list | tuple):
|
|
tensors = RxTuple(tensors)
|
|
return _ffi_api.stack(tensors, axis) # type: ignore
|
|
|
|
|
|
def collapse_sum_like(data: Expr, collapse_target: Expr) -> Expr:
|
|
"""Return a summation of data to the shape of collapse_target.
|
|
|
|
For details, please see relax.op.collapse_sum_to.
|
|
|
|
Parameters
|
|
----------
|
|
data : relax.Expr
|
|
The input tensor.
|
|
|
|
collapse_target : relax.Expr
|
|
The tensor whose shape is the shape to collapse to.
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The result tensor after summation.
|
|
"""
|
|
return _ffi_api.collapse_sum_like(data, collapse_target) # type: ignore
|
|
|
|
|
|
def collapse_sum_to(data: Expr, shape: tuple[PrimExprLike] | Expr) -> Expr:
|
|
"""Return a summation of data to the given shape.
|
|
|
|
collapse_sum_to is intended as the backward operator of tvm.relax.op.broadcast_to and
|
|
other broadcast operators in the automatic differentiation process.
|
|
|
|
We expect that data is the result of broadcasting some tensor of the given shape in some
|
|
broadcast operation. Thus the given `shape` and `data.shape` must follow broadcast rules.
|
|
|
|
During computation, all axes of `data.shape` and `shape` are checked from right to left.
|
|
For an axis, if it follows these rules, `data` will be summed over this axis:
|
|
- the axis exists in `data.shape` but not in `shape`, or
|
|
- the axis exists in `data.shape` and equals to 1 in `shape`.
|
|
|
|
Parameters
|
|
----------
|
|
data : relax.Expr
|
|
The input tensor.
|
|
|
|
shape : Union[Tuple[PrimExprLike], relax.Expr]
|
|
The shape to collapse to.
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The result tensor of the given shape after summation.
|
|
"""
|
|
if isinstance(shape, tuple | list):
|
|
shape = ShapeExpr(shape)
|
|
return _ffi_api.collapse_sum_to(data, shape) # type: ignore
|
|
|
|
|
|
def repeat(data: Expr, repeats: int, axis: int | None = None) -> Expr:
|
|
"""Repeats elements of an array.
|
|
|
|
Parameters
|
|
----------
|
|
data : relax.Expr
|
|
The input tensor.
|
|
|
|
repeats : int
|
|
The number of repetitions.
|
|
|
|
axis: Optional[int]
|
|
The axis along which to repeat values. The negative numbers are interpreted
|
|
counting from the backward. By default, use the flattened input array, and
|
|
return a flat output array.
|
|
|
|
Returns
|
|
-------
|
|
ret : relax.Expr
|
|
The computed result.
|
|
|
|
Examples
|
|
--------
|
|
.. code-block:: python
|
|
|
|
x = R.const([[1, 2], [3, 4]])
|
|
lv1 = R.repeat(x, repeats=2) # lv1 == [1, 1, 2, 2, 3, 3, 4, 4]
|
|
lv2 = R.repeat(x, repeats=2, axis=1) # lv2 == [[1., 1., 2., 2.],
|
|
# [3., 3., 4., 4.]]
|
|
"""
|
|
return _ffi_api.repeat(data, repeats, axis) # type: ignore
|
|
|
|
|
|
def tile(data: Expr, repeats: int | tuple[int] | list[int]) -> Expr:
|
|
"""Construct an array by repeating data the number of times given by repeats.
|
|
|
|
If repeats has length l, and data has dimension d, the result will have dimension of max(l, d).
|
|
|
|
If d < l, data is promoted to be l-dimensional by prepending new axes. So a shape (3,) Tensor is
|
|
promoted to (1, 3) for 2-D replication, or shape (1, 1, 3) for 3-D replication. If this is not
|
|
the desired behavior, promote data to d-dimensions manually before calling this function.
|
|
|
|
If d > l, reps is promoted to length d by pre-pending 1's to it. Thus for a data of shape
|
|
(2, 3, 4, 5), a reps of (2, 2) is treated as (1, 1, 2, 2).
|
|
|
|
Parameters
|
|
----------
|
|
data : relax.Expr
|
|
The input data to the operator.
|
|
|
|
repeats : Union[int, Tuple[int], List[int]]
|
|
The number of repetitions of data along each axis.
|
|
|
|
Returns
|
|
-------
|
|
ret : relax.Expr
|
|
The computed result.
|
|
|
|
Examples
|
|
--------
|
|
.. code-block:: python
|
|
|
|
x = R.const([[1, 2], [3, 4]])
|
|
lv1 = R.tile(x, reps=(2, 3)) # lv1 = [[1., 2., 1., 2., 1., 2.],
|
|
# [3., 4., 3., 4., 3., 4.],
|
|
# [1., 2., 1., 2., 1., 2.],
|
|
# [3., 4., 3., 4., 3., 4.]]
|
|
lv2 = R.tile(x, reps=2) # lv2 = [[1., 2., 1., 2.],
|
|
# [3., 4., 3., 4.]]
|
|
"""
|
|
if isinstance(repeats, int):
|
|
repeats = [repeats]
|
|
return _ffi_api.tile(data, repeats) # type: ignore
|
|
|
|
|
|
def flip(data, axis):
|
|
"""Reverses the order of elements along given axis while preserving array shape.
|
|
|
|
Parameters
|
|
----------
|
|
data : relax.Expr
|
|
The input data to the operator.
|
|
|
|
axis: int
|
|
The axis along which to flip over.
|
|
|
|
Returns
|
|
-------
|
|
ret : relax.Expr
|
|
The computed result.
|
|
|
|
Examples
|
|
--------
|
|
.. code-block:: python
|
|
|
|
x = [[1., 2.], [3., 4.]]
|
|
relax.flip(x, axis=0) = [[3., 4.], [1., 2.]]
|
|
|
|
relax.flip(x, axis=1) = [[2., 1.], [4., 3.]]
|
|
"""
|
|
return _ffi_api.flip(data, axis) # type: ignore
|
|
|
|
|
|
def reverse_sequence(data: Expr, seq_lengths: Expr, seq_axis: int = 1, batch_axis: int = 0) -> Expr:
|
|
"""Reverses variable length slices.
|
|
|
|
Parameters
|
|
----------
|
|
data : relax.Expr
|
|
The input tensor.
|
|
|
|
seq_lengths : relax.Expr
|
|
A 1-D tensor containing sequence lengths for each batch.
|
|
|
|
seq_axis : int
|
|
The axis along which to reverse variable length slices.
|
|
|
|
batch_axis : int
|
|
The axis that indexes the batch.
|
|
|
|
Returns
|
|
-------
|
|
ret : relax.Expr
|
|
The computed result.
|
|
"""
|
|
return _ffi_api.reverse_sequence(data, seq_lengths, seq_axis, batch_axis) # type: ignore
|
|
|
|
|
|
def gather_elements(data: Expr, indices: Expr, axis: int = 0) -> Expr:
|
|
"""Gather elements from data according to indices along the specified axis.
|
|
|
|
Parameters
|
|
----------
|
|
data : relax.Expr
|
|
The input data to the operator.
|
|
|
|
indices : relax.Expr
|
|
The indices tensor, must have integer type.
|
|
|
|
axis : int
|
|
The axis along which to index. Default is 0.
|
|
|
|
Returns
|
|
-------
|
|
ret : relax.Expr
|
|
The computed result.
|
|
|
|
Examples
|
|
--------
|
|
.. code-block:: python
|
|
|
|
data = [[1, 2], [3, 4]]
|
|
indices = [[0, 0], [1, 0]]
|
|
axis = 1
|
|
output = [[1, 1], [4, 3]]
|
|
|
|
data = [[1, 2, 3], [4, 5, 6]]
|
|
indices = [[1, 1, 1]]
|
|
axis = 0
|
|
output = [[4, 5, 6]]
|
|
"""
|
|
return _ffi_api.gather_elements(data, indices, axis) # type: ignore
|
|
|
|
|
|
def gather_nd(data: Expr, indices: Expr, batch_dims: int = 0) -> Expr:
|
|
"""Update data at positions defined by indices with values in updates.
|
|
|
|
Parameters
|
|
----------
|
|
data : relax.Expr
|
|
The input data to the operator.
|
|
|
|
indices : relax.Expr
|
|
The indices tensor, must have integer type.
|
|
|
|
batch_dims : int
|
|
The number of batch dimensions. Default is 0.
|
|
|
|
Returns
|
|
-------
|
|
ret : relax.Expr
|
|
The computed result.
|
|
|
|
Examples
|
|
--------
|
|
.. code-block:: python
|
|
|
|
batch_dims = 0
|
|
data = [[0,1],[2,3]] # data_shape = [2, 2]
|
|
indices = [[0,0],[1,1]] # indices_shape = [2, 2]
|
|
output = [0,3] # output_shape = [2]
|
|
|
|
batch_dims = 1
|
|
data = [[[0,1],[2,3]],[[4,5],[6,7]]] # data_shape = [2, 2, 2]
|
|
indices = [[1],[0]] # indices_shape = [2, 1]
|
|
output = [[2,3],[4,5]] # output_shape = [2, 2]
|
|
|
|
"""
|
|
return _ffi_api.gather_nd(data, indices, batch_dims) # type: ignore
|
|
|
|
|
|
def index_tensor(data: Expr, indices: Expr | list[Expr]) -> Expr:
|
|
"""Advanced-tensor indexing (NumPy/PyTorch-style).
|
|
|
|
Given k index tensors ``indices = (I0, I1, …, Ik-1)`` this
|
|
operator selects elements from ``data`` as if one had written
|
|
``data[I0, I1, …, Ik-1]`` in NumPy/PyTorch:
|
|
|
|
All index tensors must have an integer dtype.
|
|
|
|
Their shapes are broadcast together to a common shape ``B`` in
|
|
the usual NumPy way.
|
|
|
|
The result shape is ``B + data.shape[k:]`` (i.e. the broadcast
|
|
shape followed by the remaining axes of ``data`` that are *not*
|
|
indexed).
|
|
|
|
At compile-time Relax checks that the number of index tensors
|
|
``k`` does not exceed ``data.ndim``, that the dtypes are integer,
|
|
and that the shapes are consitent (broadcast-compatible).
|
|
|
|
Parameters
|
|
----------
|
|
data : relax.Expr
|
|
The input tensor to be indexed.
|
|
|
|
indices : Union[relax.Expr, List[relax.Expr]]
|
|
A Tuple expression containing the index tensors,
|
|
or a Python ``list`` / ``tuple`` that will be promoted to a
|
|
tuple expression automatically. Each tensor must have an
|
|
integer dtype.
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The tensor obtained after advanced indexing. Its dtype equals
|
|
``data.dtype``
|
|
|
|
Examples
|
|
--------
|
|
.. code-block:: python
|
|
|
|
import numpy as np
|
|
import tvm.relax as R
|
|
|
|
x = R.const(np.arange(9).reshape(3, 3).astype("float32"))
|
|
row = R.const(np.array([0, 2])) # shape (2,)
|
|
col = R.const(np.array([1, 0])) # shape (2,)
|
|
|
|
y = R.index_tensor(x, [row, col])
|
|
# y.shape == (2,) ; y == [1., 6.]
|
|
|
|
# Broadcasting: row : (2,1), col : (1,3) → B = (2,3)
|
|
row = R.const(np.array([[0],[1]]))
|
|
col = R.const(np.array([[0,1,2]]))
|
|
z = R.index_tensor(x, [row, col])
|
|
# z.shape == (2,3)
|
|
|
|
"""
|
|
if isinstance(indices, list | tuple):
|
|
indices = RxTuple(indices)
|
|
return _ffi_api.index_tensor(data, indices) # type: ignore
|
|
|
|
|
|
def index_put(
|
|
data: Expr,
|
|
indices: Expr | tuple[Expr],
|
|
values: Expr,
|
|
accumulate: bool = False,
|
|
) -> Expr:
|
|
"""This operation updates values in `data` at positions
|
|
specified by `indices` with corresponding values from `values`. The `indices` is a tuple
|
|
of tensors where each tensor corresponds to a dimension in `data`.
|
|
When `accumulate` is True, the operation performs accumulation (addition) rather than
|
|
replacement. The `reduction` parameter allows specifying different reduction operations.
|
|
Parameters
|
|
----------
|
|
data : relax.Expr
|
|
The input tensor to be modified
|
|
indices : Union[Expr, Tuple[Expr]]
|
|
Tuple of index tensors (one for each dimension) specifying positions to update
|
|
values : relax.Expr
|
|
Values to place at the specified indices
|
|
accumulate : bool
|
|
Whether to accumulate (add) values rather than replace (default: False)
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
A new tensor with the same shape as data but with specified positions updated
|
|
Examples
|
|
--------
|
|
.. code-block:: python
|
|
|
|
# inputs
|
|
data = torch.zeros(3, 3)
|
|
indices = (torch.tensor([0, 2]), torch.tensor([1, 1]))
|
|
values = torch.tensor([1.0, 2.0])
|
|
# output
|
|
output = [
|
|
[0.0, 1.0, 0.0],
|
|
[0.0, 0.0, 0.0],
|
|
[0.0, 2.0, 0.0],
|
|
]
|
|
# with accumulate=True
|
|
output = [
|
|
[0.0, 1.0, 0.0],
|
|
[0.0, 0.0, 0.0],
|
|
[0.0, 3.0, 0.0],
|
|
]
|
|
"""
|
|
if isinstance(indices, list | tuple):
|
|
indices = RxTuple(indices)
|
|
return _ffi_api.index_put(data, indices, values, accumulate) # type: ignore
|
|
|
|
|
|
def meshgrid(tensors: Expr | list[Expr], indexing: str | None = "ij") -> Expr:
|
|
"""Generate coordinate grids from input tensors.
|
|
|
|
Parameters
|
|
----------
|
|
tensors : Union[relax.Expr, List[relax.Expr]]
|
|
An Expr in Tuple type, containing 1D tensors (or scalars promoted to 1D)
|
|
to generate coordinate grids from, or a list of such tensors.
|
|
|
|
indexing : Optional[str]
|
|
The indexing mode, either "ij" (matrix indexing) or "xy" (Cartesian indexing).
|
|
Defaults to "ij".
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
A Tuple of tensors representing the coordinate grids.
|
|
"""
|
|
if isinstance(tensors, list | tuple):
|
|
tensors = RxTuple(tensors)
|
|
return _ffi_api.meshgrid(tensors, indexing)
|
|
|
|
|
|
def scatter_elements(
|
|
data: Expr, indices: Expr, updates: Expr, axis: int = 0, reduction: str = "update"
|
|
):
|
|
"""ONNX style scatter elements. This operation updates its value in `data` to values
|
|
specified by `updates` at specific index positions specified by `indices`.
|
|
For example, in 2D tensor, the update corresponding to the [i][j] entry is performed
|
|
as below:
|
|
|
|
.. code-block::
|
|
|
|
output[indices[i][j]][j] = updates[i][j] if axis = 0
|
|
output[i][indices[i][j]] = updates[i][j] if axis = 1
|
|
|
|
When the `reduction` is set to some reduction function `f`, the update corresponding to
|
|
[i][j] entry is performed as below:
|
|
|
|
.. code-block::
|
|
|
|
output[indices[i][j]][j] += f(output[indices[i][j]][j], updates[i][j]) if axis = 0
|
|
output[i][indices[i][j]] += f(output[i][indices[i][j]], updates[i][j]) if axis = 1
|
|
|
|
Where `f` is update, add, mul, mean, max, min.
|
|
|
|
Parameters
|
|
----------
|
|
data : relax.Expr
|
|
The input data to the operator.
|
|
|
|
indices: relax.Expr
|
|
The index positions to update in `data`.
|
|
|
|
updates: relax.Expr
|
|
Values to replace to.
|
|
|
|
axis: int
|
|
Axis to scatter on.
|
|
|
|
reduction: str
|
|
Type of reduction to apply: update, add, mul, mean, max, min.
|
|
It is "update" by default.
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The result has the same size as data, and the same shape as data
|
|
|
|
Examples
|
|
--------
|
|
.. code-block:: python
|
|
|
|
# inputs
|
|
data = [
|
|
[0.0, 0.0, 0.0],
|
|
[0.0, 0.0, 0.0],
|
|
[0.0, 0.0, 0.0],
|
|
]
|
|
indices = [
|
|
[1, 0, 2],
|
|
[0, 2, 1],
|
|
]
|
|
updates = [
|
|
[1.0, 1.1, 1.2],
|
|
[2.0, 2.1, 2.2],
|
|
]
|
|
axis = 0
|
|
reduction = "update"
|
|
|
|
# output P
|
|
output = [
|
|
[2.0, 1.1, 0.0]
|
|
[1.0, 0.0, 2.2]
|
|
[0.0, 2.1, 1.2]
|
|
]
|
|
|
|
"""
|
|
return _ffi_api.scatter_elements(data, indices, updates, axis, reduction) # type: ignore
|
|
|
|
|
|
def scatter_nd(data: Expr, indices: Expr, updates: Expr, reduction: str = "update") -> Expr:
|
|
"""Scatter updates into an array according to indices.
|
|
|
|
Parameters
|
|
----------
|
|
data: relax.Expr
|
|
The input data to be updated.
|
|
|
|
indices: relax.Expr
|
|
The index positions to update in `data`.
|
|
|
|
updates: relax.Expr
|
|
Values to replace to.
|
|
|
|
reduction: str
|
|
Type of reduction to apply: update, add, mul, max, min.
|
|
It is "update" by default.
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The result has the same shape as data.
|
|
|
|
Examples
|
|
--------
|
|
.. code-block:: python
|
|
|
|
# inputs
|
|
data = [1, 2, 3, 4, 5, 6, 7, 8]
|
|
indices = [[4], [3], [1], [7]]
|
|
updates = [9, 10, 11, 12]
|
|
|
|
# output
|
|
output = [1, 11, 3, 10, 9, 6, 7, 12]
|
|
|
|
"""
|
|
return _ffi_api.scatter_nd(data, indices, updates, reduction) # type: ignore
|
|
|
|
|
|
def slice_scatter(input_tensor: Expr, src: Expr, start, end, step, axis=0):
|
|
"""Embeds the values of the src tensor into input at the given dimension.
|
|
|
|
Parameters
|
|
----------
|
|
input_tensor: relax.Expr
|
|
The input tensor to be updated.
|
|
|
|
src: relax.Expr
|
|
The tensor to embed into input.
|
|
|
|
axis: int
|
|
The dimension to insert the slice into.
|
|
|
|
start:
|
|
The start index of where to insert the slice.
|
|
|
|
end:
|
|
The end index of where to insert the slice.
|
|
|
|
step:
|
|
The how many elements to skip in.
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The computed result tensor with the same shape as `data`.
|
|
|
|
"""
|
|
if not is_prim_expr(start):
|
|
start = prim_value(start)
|
|
if not is_prim_expr(end):
|
|
end = prim_value(end)
|
|
if not is_prim_expr(step):
|
|
step = prim_value(step)
|
|
return _ffi_api.slice_scatter(input_tensor, src, axis, start, end, step)
|
|
|
|
|
|
def one_hot(
|
|
indices: Expr,
|
|
on_value: int | float | Expr,
|
|
off_value: int | float | Expr,
|
|
depth: int,
|
|
axis: int = -1,
|
|
) -> Expr:
|
|
"""Returns a one-hot tensor.
|
|
|
|
Parameters
|
|
----------
|
|
indices : relax.Expr
|
|
The indices to set to `on_value`.
|
|
|
|
on_value : int | float | Expr
|
|
The value to fill at `indices`.
|
|
|
|
off_value : int | float | Expr
|
|
The value to fill at other locations.
|
|
|
|
depth : int
|
|
The depth of the one-hot dimension.
|
|
|
|
axis : int, optional
|
|
The axis to fill. Default is -1 which adds a new dimension at the end.
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The computed result.
|
|
|
|
Examples
|
|
--------
|
|
.. code-block:: python
|
|
|
|
indices = [0, 1, 2]
|
|
depth = 3
|
|
on_value = 1
|
|
off_value = 0
|
|
|
|
one_hot(indices, on_value, off_value, depth) =
|
|
[[1, 0, 0],
|
|
[0, 1, 0],
|
|
[0, 0, 1]]
|
|
"""
|
|
on_value = prim_value(on_value)
|
|
off_value = prim_value(off_value)
|
|
return _ffi_api.one_hot(indices, on_value, off_value, depth, axis) # type: ignore
|