GatherOperation

Applicable Products

Hardware Model	Supported	Remarks
Atlas 350 accelerator cards	√	-
Atlas A3 inference products/Atlas A3 training products	√	-
Atlas A2 training products/Atlas A2 inference products	√	-
Atlas training products	√	-
Atlas inference products	√	-
Atlas 200I/500 A2 inference products	√	The input x data type can only be float16.

Description

Collects slices from the input tensor based on the index and combines these slices into a new tensor.

Figure 1 Context of the GatherOperation operator

Definition

struct GatherParam {
    int64_t axis = 0;
    int64_t batchDims = 0;
    uint8_t rsv[16] = {0};
};

Parameters

Member	Type	Default Value	Value Range	Mandatory	Description
axis	int64_t	0	≥ 0 and ≤ n	Yes	Axis from which slices are to be collected. Defaults to 0. n is the maximum index of the dimension of input x.
batchDims	int64_t	0	≥ 0 and ≤ min(axis, m)	Yes	Number of batch processing dimensions. The parameter allows you to gather data slices from each element of a batch. For example, if batchDims is 1, there is an outer loop on the (axis – batchDims) axis of x. For details, see Example 2. m is the maximum index of the dimension of input indices.
rsv[16]	uint8_t	{0}	[0]	No	Reserved

Member

Type

Default Value

Value Range

Mandatory

Description

axis

int64_t

≥ 0 and ≤ n

Yes

Axis from which slices are to be collected. Defaults to 0. n is the maximum index of the dimension of input x.

batchDims

int64_t

≥ 0 and ≤ min(axis, m)

Yes

Number of batch processing dimensions. The parameter allows you to gather data slices from each element of a batch.

For example, if batchDims is 1, there is an outer loop on the (axis – batchDims) axis of x. For details, see Example 2.

m is the maximum index of the dimension of input indices.

rsv[16]

uint8_t

{0}

[0]

Reserved

Input

Parameter	Dimension	Data Type	Format	Description
x	[dim_x_0, dim_x_1, ..., dim_x_n]	float16/float/bfloat16/int32/uint32/int64	ND	Input tensor.
indices	When batchDims = 0: [dim_in_0, dim_in_1, ..., dim_in_m] When batchDims > 0: [dim_x_0, ..., dim_x_(batchDims-1), dim_in_k, ..., dim_in_m]	int64/int32/uint32	ND	Index table. The element value range is [0, dim_x_axis).

Output

Parameter	Dimension	Data Type	Format	Description
output	[dim_x_0, dim_x_1, ..., dim_x_n]	float16/float/bfloat16/int32/uint32/int64	ND	Output tensor. The data type is the same as that of x. If batchDims is 0, dim_x_axis in the output dimension is replaced with dim_in_0,dim_in_1,...,dim_in_m. If batchDim is greater than 0, dim_x_axis in the output dimension is replaced with dim_in_k,...,dim_in_m.

Parameter

Dimension

Data Type

Format

Description

output

[dim_x_0, dim_x_1, ..., dim_x_n]

float16/float/bfloat16/int32/uint32/int64

Output tensor. The data type is the same as that of x.

If batchDims is 0, dim_x_axis in the output dimension is replaced with dim_in_0,dim_in_1,...,dim_in_m.
If batchDim is greater than 0, dim_x_axis in the output dimension is replaced with dim_in_k,...,dim_in_m.

Constraints

The number of dimensions of output is less than or equal to 8. Therefore, (n + 1) + (m + 1) -1-batchdims <= 8.

API Calling Example

Example 1:

axis = 1;
batchDims = 0;
The input tensor is as follows:
 x=  [[1,2,3],
      [4,5,6],
      [7,8,9]]
The indices tensor is as follows:
indices=[2,0]
Obtain the x data slice on the axis based on the value of the indices tensor. The output tensor is as follows:
output=[[3, 1],
        [6, 4],
        [9, 7]]

Example 2:

axis= 1;
batchDims = 1;
The input tensor is as follows:
 x=  [[1,2,3],
      [4,5,6],
      [7,8,9]]
The indices tensor is as follows:
indices=[[1],
         [2],
         [0]]
batch_dims = 1 indicates that batch processing is performed on the first axis (axis 0). On axis 0, the corresponding gather processing is performed on x[i] and indices[i]. The data slice of x is obtained on the axis based on the value of indices. i is the coordinate of the batch axis. The output tensor is as follows:
output= [[2,],
         [6,],
         [7,]]
The above is equivalent to:
def gather_fun(x, indices, axis):
       batch_dims=1
       res= []
    # Perform an outer loop.
    for p,i in zip(x, indices): 
         r = tf.gather(p, i, axis=axis-batch_dims)
         res.append(r)
    return tf.stack(res)

Parent topic: atb/infer_op_params.h