GatherOperation
Description
Collects slices from the input tensor based on the index and combines these slices into a new tensor.
Definition
struct GatherParam {
int64_t axis = 0;
int64_t batchDims = 0;
uint8_t rsv[16] = {0};
};
Parameters
Member |
Type |
Default Value |
Description |
|---|---|---|---|
axis |
int64_t |
0 |
Axis from which slices are to be collected. The default value is 0. The value of axis must be greater than or equal to 0. |
batchDims |
int64_t |
0 |
Allowed different items to be collected from each element of a batch. The default value is 0. The value of batchDims must be greater than or equal to 0 and less than or equal to axis. |
rsv[16] |
uint8_t |
{0} |
Reserved |
Input
Parameter |
Dimension |
Data Type |
Format |
Description |
|---|---|---|---|---|
x |
[-1,...,-1] The value -1 indicates that the size of the current dimension is not restricted. |
float16/float/bf16/int32/uint32 |
ND |
Input tensor. |
indexs |
[-1,...,-1] The value -1 indicates that the size of the current dimension is not restricted. |
int64/int32/uint32 |
ND |
Index table. The value must be in the range of [0, x.shape[axis]]. The sum of the number of dimensions of x and indexes must be less than or equal to 9. The number of dimensions of indexs must be greater than or equal to batchdims. |
Output
Parameter |
Dimension |
Data Type |
Format |
Description |
|---|---|---|---|---|
output |
[-1,...,-1] The value -1 indicates that the size of the current dimension is not restricted. |
float16/float/bf16/int32/uint32 |
ND |
Output tensor. The data type is the same as that of x. |
Restrictions
- The value of the indices index table must be within the range of [0, x.shape[axis]]. The sum of the number of dimensions of x and indexes must be less than or equal to 9. The number of dimensions of indexs must be greater than or equal to batchdims.
- The value of axis must be greater than or equal to 0.
- The value of batchDims must be greater than or equal to 0 and less than or equal to axis.
API Calling Example
axis = 0;
batchDims = 0;
>>> x
tensor([[3, 7, 1, 8],
[2, 6, 5, 0],
[1, 4, 6, 9]])
>>> indices
tensor([[1, 0],
[0, 2],
[2, 1]])
>>> output
tensor([[[2, 6, 5, 0],
[3, 7, 1, 8]],
[[3, 7, 1, 8],
[1, 4, 6, 9]],
[[1, 4, 6, 9],
[2, 6, 5, 0]]])
axis= 1;
batchDims = 1;
>>> output
tensor([[7, 3],
[2, 5],
[6, 4]])