EmbeddingBag

Description

Obtains a group of numbers from weight based on indices, and performs aggregation in max, sum, or mean mode on the obtained numbers based on offsets and the aggregation mode specified by mode. Other parameters refine the control of the computation process.

Input

• weight: 2-dimensional non-contiguous input tensor of type float or shape. It is a word embedding matrix. The data type can be ND.
• indices: 1- or 2-dimensional input tensor of type UINT8, INT8, INT16, INT32, INT64, or shape. It is a tensor that contains indices and specifies the embedding vectors of the words to be extracted from weight.
• offset (optional): input tensor, which is used to split indices into multiple bags. When indices is 1-dimensional, the shape of offsets can be 1-dimensional. When indices is 2-dimensional, offsets can be 1- or 2-dimensional. The data type can be UINT8, INT8, INT16, INT32, or INT64.
• per_sample_weights (optional): sample weight. The shape can be 1-dimensional. The data type is the same as that of weight. The value can be nullptr only in sum mode.

Attribute

• mode: string, specifying the aggregation mode. 0: sum mode; 1: mean mode; other values: max mode.
• scale_grad_by_freq: input tensor of type bool, which controls whether to scale the gradient based on the word frequency. When scale_grad_by_freq is true, the gradient is scaled based on the word frequency. When scale_grad_by_freq is false, the gradient is not scaled.
• sparse: bool, which is used to control the sparse mode. If the value is false, it indicates that the weight is not a sparse matrix. If the value is true, it indicates that the weight is a sparse matrix.
• include_last_offset: bool, specifying whether the last offset is included. If the value is false, the last offset is not included. If the value is true, the last offset is included.

Output

y: output tensor of type float.

Constraints

None.

ONNX Opset Support

Opset v8/v9/v10/v11/v12/v13/v14/v15/v16